公式
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;
}
投稿日時:
最終更新: