Hard

Sum of Distances in Tree

LAB

Given a tree with n nodes, return for every node the sum of distances to all other nodes.

EXAMPLES

Example 1

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

Output
[
  8,
  12,
  6,
  10,
  10,
  10
]

FUNCTION SHAPE

n: intedges: intMatrix→intArray

SOLUTION NOTE

Re-rooting technique. First DFS computes answer for root. Second DFS propagates to children using relationship between parent/child answers.

Reveal reference solution +

pythonREFERENCE

def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
    graph = defaultdict(list)
    for u, v in edges:
        graph[u].append(v)
        graph[v].append(u)

    count = [1] * n  # Subtree size
    res = [0] * n

    # First DFS: compute count and res[0]
    def dfs1(node, parent):
        for child in graph[node]:
            if child != parent:
                dfs1(child, node)
                count[node] += count[child]
                res[0] += res[child] + count[child]

    # Second DFS: compute res for all nodes
    def dfs2(node, parent):
        for child in graph[node]:
            if child != parent:
                # Moving from node to child:
                # count[child] nodes get closer by 1
                # n - count[child] nodes get farther by 1
                res[child] = res[node] - count[child] + (n - count[child])
                dfs2(child, node)

    dfs1(0, -1)
    dfs2(0, -1)
    return res

TimeO(n)

SpaceO(n)

Open on LeetCode ↗