Graph Theory37 sections · 1633 units
Open in Course

Tree Distances II - Implementation (Two DFS passes)

(C++ code)

Here is the two-pass DFS implementation:

function treeDistancesII(n, adj):
    subtreeSize = array of size n + 1, all 0
    subtreeSum = array of size n + 1, all 0
    answer = array of size n + 1, all 0

    // DFS 1: compute subtree sizes and sums rooted at node 1
    function dfs1(u, parent):
        subtreeSize[u] = 1
        subtreeSum[u] = 0
        for v in adj[u]:
            if v != parent:
                dfs1(v, u)
                subtreeSize[u] = subtreeSize[u] + subtreeSize[v]
                subtreeSum[u] = subtreeSum[u] + subtreeSum[v] + subtreeSize[v]

    // DFS 2: reroot to compute answer for each node
    function dfs2(u, parent):
        for v in adj[u]:
            if v != parent:
                // When rerooting from u to v:
                // Nodes in v's subtree get 1 closer
                // All other nodes get 1 farther
                answer[v] = answer[u] - subtreeSize[v] + (n - subtreeSize[v])
                dfs2(v, u)

    dfs1(1, 0)
    answer[1] = subtreeSum[1]
    dfs2(1, 0)

    print answer[1..n]

Time: O(n)O(n). Space: O(n)O(n).