公式

F - Distance Sums 2 解説 by blackyuki


\(\sum_{j=1}^{N}dis(i,j)\)\(ans(i)\) と略記します。

まず各頂点 \(i\) について \(dis(1,i)\) を求めます。これは深さ優先探索などによって合計 \(O(N)\) で計算できるので、\(ans(1)\) を求めることができます。

次に、頂点 \(1\) に隣接する頂点の \(1\) つを \(s\) として、\(ans(1)\)\(ans(s)\) の差分を考えます。ある頂点 \(j\) が頂点 \(s\) の部分木に含まれる場合 \(dis(s,j)=dis(i,j)-1\) 、含まれない場合 \(dis(s,j)=dis(i,j)+1\) なので、以下が成り立ちます。

\(ans(s) = ans(1) + N - 2 \times sub(s)\) (ただし \(sub(s)\) は頂点 \(s\) の部分木のサイズ)

よってあらかじめ各頂点について部分木のサイズを計算した後、頂点 \(1\) に近い方から \(ans(i)\) を求めることで、全体で \(O(N)\) で解けました。

実装例 (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;
}

投稿日時:
最終更新: