Graph Theory37 sections · 1633 units
Open in Course

Tree Distances II - Implementation

(Rerooting DP)

Here is the solution:

sub = array of size n
down = array of size n
ans = array of size n

function dfs1(u, p):
    sub[u] = 1
    for v in adj[u]:
        if v == p:
            continue
        dfs1(v, u)
        sub[u] = sub[u] + sub[v]
        down[u] = down[u] + down[v] + sub[v]

function dfs2(u, p, up):
    ans[u] = down[u] + up
    for v in adj[u]:
        if v == p:
            continue
        dfs2(v, u, up + (down[u] - down[v] - sub[v]) + (n - sub[v]))

This runs in O(n)O(n) time and uses O(n)O(n) space.