CSES - Tree Distances II

Time Complexity: $\mathcal{O}(n)$

It's easy to find the sum of distances from a single node - just root the tree at that node, do a DFS, and add the depths of each other node to the answer. Unfortunately, $n$ can go up to $2\cdot 10^5$, so we can't just do this for each node.

If we have the answer for some node (let's say node 1), how can we quickly find the answer for its neighbours?

The key observation is that if we reroot the tree at node 1's neighbour (let's say node 2), then

• The depths of all nodes in node 2's subtree decrease by 1.
• The depths of all nodes outside of its subtree increase by 1.

This gives us a nice way to transition from node 1's answer to node 2's answer using only $n$ and the size of node 2's subtree! Observe that the change in the answer is exactly $n-2(\text{node 2’s subtree size})$.

We can use DFS to find both the answer for node 1 and the size of each node's subtree when rooted at node 1, and then DFS again to compute all the answers.

1#include <bits/stdc++.h>
2typedef long long ll;
3using namespace std;
4
5int n;
6vector<int> graph[200001];
7ll dp[200001], ans[200001];
8
9void dfs1(int node = 1, int parent = 0, ll depth = 0) {
10    ans[1] += depth;