Hard

Minimum Weighted Subgraph With the Required Paths II

LAB

Given directed weighted edges and sources src1, src2, and destination, return the minimum total edge weight of a subgraph allowing both sources to reach destination, or -1.

EXAMPLES

Example 1

Input
{
  "n": 6,
  "edges": [
    [
      0,
      2,
      2
    ],
    [
      0,
      5,
      6
    ],
    [
      1,
      0,
      3
    ],
    [
      1,
      4,
      5
    ],
    [
      2,
      1,
      1
    ],
    [
      2,
      3,
      3
    ],
    [
      2,
      3,
      4
    ],
    [
      3,
      4,
      2
    ],
    [
      4,
      5,
      1
    ]
  ],
  "src1": 0,
  "src2": 1,
  "dest": 5
}

Output
9

FUNCTION SHAPE

n: intedges: intMatrixsrc1: intsrc2: intdest: int→int

SOLUTION NOTE

Binary lifting for LCA. Min subtree connecting 3 nodes meets at the deepest pairwise LCA.

Reveal reference solution +

pythonREFERENCE

def minimumWeight(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
    from collections import defaultdict
    n = len(edges) + 1
    LOG = (n - 1).bit_length()

    tree = defaultdict(list)
    for u, v, w in edges:
        tree[u].append((v, w))
        tree[v].append((u, w))

    parent = [[-1] * n for _ in range(LOG)]
    depth, dist = [0] * n, [0] * n

    def dfs(u, p):
        for v, w in tree[u]:
            if v != p:
                parent[0][v] = u
                depth[v] = depth[u] + 1
                dist[v] = dist[u] + w
                dfs(v, u)

    dfs(0, -1)
    for k in range(1, LOG):
        for v in range(n):
            if parent[k-1][v] != -1:
                parent[k][v] = parent[k-1][parent[k-1][v]]

    def lca(u, v):
        if depth[u] < depth[v]: u, v = v, u
        diff = depth[u] - depth[v]
        for k in range(LOG):
            if diff & (1 << k): u = parent[k][u]
        if u == v: return u
        for k in range(LOG - 1, -1, -1):
            if parent[k][u] != parent[k][v]:
                u, v = parent[k][u], parent[k][v]
        return parent[0][u]

    def path_dist(u, v):
        return dist[u] + dist[v] - 2 * dist[lca(u, v)]

    res = []
    for s1, s2, d in queries:
        l12, l1d, l2d = lca(s1, s2), lca(s1, d), lca(s2, d)
        meet = max([l12, l1d, l2d], key=lambda x: depth[x])
        res.append(path_dist(s1, meet) + path_dist(s2, meet) + path_dist(d, meet))
    return res

TimeO(n log n + q log n)

SpaceO(n log n)

Open on LeetCode ↗