Approach #1: Subtree Sum and Count [Accepted]
Intuition
Let ans
be the returned answer, so that in particular ans[x]
be the answer for node x
.
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[x]
and 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)
.
Algorithm
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 count
and stsum
using a postorder traversal, where on exiting some node
, the count
and 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 root
: 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 child
than parent
, and Ncount[child]
nodes that are 1
harder to get to from child
than parent
.
Using a second, preorder traversal, we can update our answer in linear time for all of our nodes.
Complexity Analysis

Time Complexity: , where is the number of nodes in the graph.

Space Complexity: .
Analysis written by: @awice.