Approach #1: Subtree Sum and Count [Accepted]
ans be the returned answer, so that in particular
ans[x] be the answer for node
Naively, finding each
ans[x] would take time (where is the number of nodes in the graph), which is too slow. This is the motivation to find out how
ans[y] are related, so that we cut down on repeated work.
Let's investigate the answers of neighboring nodes and . In particular, say is an edge of the graph, that if cut would form two trees (containing ) and (containing ).
Then, as illustrated in the diagram, the answer for in the entire tree, is the answer of on
"x@X", plus the answer of on
"y@Y", plus the number of nodes in
"#(Y)". The last part
"#(Y)" is specifically because for any node
z in Y,
dist(x, z) = dist(y, z) + 1.
By similar reasoning, the answer for in the entire tree is
ans[y] = x@X + y@Y + #(X). Hence, for neighboring nodes and ,
ans[x] - ans[y] = #(Y) - #(X).
Root the tree. For each node, consider the subtree of that node plus all descendants. Let
count[node] be the number of nodes in , and
stsum[node] ("subtree sum") be the sum of the distances from
node to the nodes in .
We can calculate
stsum using a post-order traversal, where on exiting some
stsum of all descendants of this node is correct, and we now calculate
count[node] += count[child] and
stsum[node] += stsum[child] + count[child].
This will give us the right answer for the
ans[root] = stsum[root].
Now, to use the insight explained previously: if we have a node
parent and it's child
child, then these are neighboring nodes, and so
ans[child] = ans[parent] - count[child] + (N - count[child]). This is because there are
count[child] nodes that are
1 easier to get to from
N-count[child] nodes that are
1 harder to get to from
Using a second, pre-order traversal, we can update our answer in linear time for all of our nodes.
Time Complexity: , where is the number of nodes in the graph.
Space Complexity: .
Analysis written by: @awice.