ARC086 E : Smuggling Marbles(1000)
作成日：2017.12.12
最終更新日：2017.12.12
問題概要
各頂点に$0$から$N$までの番号が付いた頂点$0$を根とする$N+1$頂点の根付き木がある。
頂点$i$$(i \geq 1)$の親は頂点$p_i$である。
頂点のいくつかに玉を置いて以下の手順をいずれかの頂点に玉がある限り続ける。
- 頂点$0$に玉が置かれているならば、その球を箱に移す。
- すべての玉を今ある頂点の親の頂点に同時に移す。
- 2つ以上の玉が置かれている頂点にある玉をすべて取り除く。
手順を終えた後に箱にある玉の数を玉の初期配置それぞれについて求め、その和を$10^9 + 7$で割った余りを答えよ。
制約
$1 \leq N \leq 2 \leq 10^5$
$0 \leq p_i < i$
考え方
深さが異なる頂点に置かれた玉が同時に同じ頂点に辿り着くことはないので、初めに玉がある頂点の深さごとに考えれば良い。
また場合の数を考えるのは煩雑なので、玉がある確率を考えることにする。
この時最終的な答えは、深さごとに求めた箱に玉が入る確率の総和に$2^{N+1}$をかけたものである。
ある深さにある玉が箱に入る確率を考える。
ある深さにある頂点が1つしかない時は明らかに$1/2$である。
そうでない時は玉が頂点に2つ以上置かれることを考慮する必要がある。
ある頂点$v$の子に現在玉が置かれている確率が分かっている時に、1回手順を行った後に頂点$v$に玉が置かれている確率を求めたい。
これは子を順に見ていき、玉が0個置かれる確率$p_0$と1個置かれる確率$p_1$を持っていけば良い。
具体的には初め$p_0 = 1$,$p_1 = 0$とし、玉が置かれている確率が$q$である子それぞれに対し$p_1 = p_1 + p_0 q$ , $p_0 = p_0 (1-q)$と更新していけば良い。
これを愚直にやると$O(N^2)$となり間に合わないので工夫する。
まずは深さごとに何度も木DPを行うのではなく、木DPを1度だけ行うようにしなければならない。
それには各頂点ごとに手順を行った回数ごとに玉が置かれている確率を配列で持つようにすれば良い。
このようにした時、先の計算は2つの配列をマージする操作となることが分かる。
よって、より大きな配列に小さい配列の情報を組み入れる「マージする時の一般的なテク」を使うことができる。
但し、先に述べた方法における$p_0$は見ている頂点の深さが変わるごとに初期化する必要があること、さらにマージが起きない位置まで初期化すると計算量が減らないことに注意する。
最後にこの方法の計算量を考える。
配列の要素は頂点を1つ見るごとに1だけしか増えないので、マージする要素の個数の総和は$O(N)$である。
この総和のオーダーが計算量のオーダーと等しいから、計算量も$O(N)$であり十分高速である。
ソースコード
マクロ等はこちら
const LL mod = 1e9 + 7;
const int MC = 200010;
int N;
vector<int> V[MC];
vector<LL> ans[MC];
vector<LL> ans0[MC];
int reset[MC];
LL inv2 = (mod+1) / 2;
void dfs(int q){
reset[q] = -1;
for(auto u : V[q]){
dfs(u);
if(ans[q].size() < ans[u].size()){
swap(ans[q],ans[u]);
swap(ans0[q],ans0[u]);
swap(reset[q],reset[u]);
}
int z = ans[q].size() - ans[u].size();
while(reset[u] >= 0){
ans0[u][reset[u]] = (mod+1-ans[u][reset[u]]) % mod;
--reset[u];
}
while(reset[q] >= z){
ans0[q][reset[q]] = (mod+1-ans[q][reset[q]]) % mod;
--reset[q];
}
repp(i,0,ans[u].size()){
ans[q][z+i] = (ans0[q][z+i] * ans[u][i] + ans0[u][i] * ans[q][z+i]) % mod;
(ans0[q][z+i] *= ans0[u][i]) %= mod;
}
}
ans[q].PB(inv2);
ans0[q].PB(inv2);
reset[q] = ans[q].size()-2;
}
int main(){
cin >> N;
repp(i,1,N+1){
int p;
cin >> p;
V[p].PB(i);
}
dfs(0);
repp(i,1,ans[0].size()) (ans[0][0] += ans[0][i]) %= mod;
repp(i,0,N+1) (ans[0][0] *= 2) %= mod;
cout << ans[0][0] << endl;
return 0;
}解法まとめ
手順を行った回数ごとに頂点に玉が置かれている確率を木DPで求める。
その時にマージテクを用いる。(10-36行)
最終的な答えは、深さ(=手順を行った回数)ごとに求めた箱に玉が入る確率の総和に$2^{N+1}$をかけたものである。(46-48行)