Official

F - Distance Sums 2 Editorial by en_translator


Let us abbreviate \(\sum_{j=1}^{N}dis(i,j)\) to \(ans(i)\).

First, we compute \(dis(1,i)\) for each vertex \(i\). This can be computed with a Depth-First Search (DFS) in a total of \(O(N)\) time, so we can obtain \(ans(1)\).

Next, for a vertex \(s\) adjacent to Vertex \(1\), we consider the difference of \(ans(1)\) and \(ans(s)\). If a vertex \(j\) is included in a subtree under Vertex \(s\), then \(dis(s,j)=dis(i,j)-1\); if not, \(dis(s,j)=dis(i,j)+1\); so the following relation holds:

\(ans(s) = ans(1) + N - 2 \times sub(s)\) (where \(sub(s)\) is the size of subtree under Vertex \(s\))

Therefore, we can precalculate the size of subtree under each vertex and compute \(ans(i)\) from the nearest to Vertex \(1\) to the furthest, so that the problem is solved in a total of \(O(N)\) time.

Sample code (C++)

#include <bits/stdc++.h>
using namespace std;

int main(){
    int n; cin >> n;
    vector<vector<int>> g(n);
    for(int i = 0; i < n - 1; i++) {
        int a, b; cin >> a >> b;
        g[a-1].push_back(b-1);
        g[b-1].push_back(a-1);
    }
    vector<long long> sub(n,1), ans(n);
    function<void(int, int, int)> dfs=[&](int i, int p, int d){
        ans[0] += d;
        for(int x : g[i]) if(x != p) {
            dfs(x, i, d+1);
            sub[i] += sub[x];
        }
    }; dfs(0, -1, 0);
    function<void(int, int)> dfs2=[&](int i, int p){
        for(int x : g[i]) if(x != p) {
            ans[x] = ans[i] - 2 * sub[x] + n;
            dfs2(x, i);
        }
    }; dfs2(0, -1);
    for(long long x : ans) cout << x << endl;
}

posted:
last update: