[ABC220] F - Distance Sums 2
[Q] https://atcoder.jp/contests/abc220/tasks/abc220_f
順位表をみたら、数分で何十人も解答者がいたので、どこかに転がっているだろうと思い検索しまくった。しかし何も得られず…。
解いてなお、数分で解けるのはすごいと思った。
1. とりあえずUnionFindを張って、頂点0を親としたときの、各頂点のsizeをとる。（UnionFindの機能にあるsize取得を使っただけ。Uniteしないからべつでも）
2. 頂点0からのans[0]だけ出す。O(N)で探索。
3. 0に隣接する頂点の解答を、ans[0]を再利用して出す。sizeの差分なので、O(1)で求められる。
2-1-0-3-4
だとして。
1. 0を親にしたとき、各頂点のサイズは？
size[0]=5
size[1]=2
size[2]=1
size[3]=2
size[4]=1
2. とりあえず頂点0の解答だけ探索する。
ans[0] = 6
3. ans[0]の差分を利用して、隣接する頂点の解答を出していく。
・0->1に移動すると、
ans[1] = ans[0] - size[1] + (5-size[1]) = 7
1側の頂点{1,2} は距離が1ずつ縮まって-2 = size[1]
0側の頂点{0,3,4} は距離が1ずつ遠のいて+3 = N-size[1]
・1->2に移動すると、
ans[2] = ans[1] - size[2] + (5-size[2]) = 10
2側の頂点{2} が縮まって-1
1側の頂点{1,0,3,4} が遠のいて+4
・0->3に移動すると、
ans[3] = ans[0] - size[3] + (5-size[3]) = 7
3側の頂点{3,4} が縮まって-2
0側の頂点{2,1,0} が遠のいて+3
・3->4に移動すると、
ans[4] = ans[3] - size[4] + (5-size[4]) = 10
4側の頂点{4} が縮まって-1
3側の頂点{2,1,0,3} が遠のいて+4
ans[] = {6,7,10,7,10} が解答。Q. UnionFind？
A. テンプレートがあるので、今実装と切り分けるのがいいかも。
頂点のサイズをとれるテンプレートを利用する。
実装。
ll ans[200010];
void dfs(ll u, ll p, ll d){ // ans[0]入れるだけ
for(auto v: G[u]){
if (v==p) continue;
ans[0]+=d;
dfs(v, u, d+1);
}
}
void dfs2(ll u, ll p, UnionFind &tree, ll &N){ // 頂点0から1個ずつ動いて、ansを埋める
for(auto v: G[u]){
if (v==p) continue;
// 親から子にうつったときに、子供側の距離が1縮まって、親側の距離が1ずつ増える
ans[v] = ans[u]-tree.siz[v] + (N-tree.siz[v]);
dfs2(v, u, tree, N);
}
}
int main(){
ll N;
cin >> N;
G.assign(N, vector<ll>());
seen.assign(N, false);
rep(i, N-1){
ll a, b;
cin >> a >> b;
--a, --b;
G[a].emplace_back(b);
G[b].emplace_back(a);
}
UnionFind tree(N);
rep(i, N){ // 親入れ
seen[i] = true;
tree.mkpar(i, i); // (now, set)
}
dfs(0, -1, 1); // ans[0]入れるだけ
dfs2(0, -1, tree, N); // 頂点0から、1個ずつ移動してはansを埋める
rep(i, N) cout << ans[i] << endl;
}