E - Select from Subtrees Editorial
by
mechanicalpenciI
リスがアメを選ぶ順番について、どのような順番にしても答えは変わらないため、今回は木 \(T\) の頂点を帰りがけ順に列挙したときの順番で、対応するリスがアメを選ぶことを考えます。
ここで、頂点 \(i\) を根とした部分木に属する頂点 \(j\) にわたる \(C_j\) の和を \(S_i\), \(D_j\) の和を \(T_i\) とします。帰りがけに頂点を並べた際、頂点 \(i\) の子孫はすべて頂点 \(i\) より前に並び、祖先はすべて頂点 \(i\) より後ろに並ぶため、リス \(i\) が選べるアメの候補はそれ以前のリスのアメの選び方によらず、\(S_i-(T_i-D_i)\) 個あり、リス \(i\) はその中から \(D_i\) 個のアメを選びます。すなわち、リス \(i\) のアメの選び方は、それ以前の他のリスのアメの選び方によらず \(\binom{S_i-T_i+D_i}{D_i}\) 通りあります。ここで、 \(S_i-T_i+D_i<D_i\) 、すなわち \(S_i<T_i\) ならばリス \(i\) は条件をみたすようにアメを選ぶことはできず、特に答えは \(0\) となります。
最終的な答えはこれを \(i=1,2,\ldots,N\) にわたりかけ合わせた値
\[ \prod_{i=1}^n \binom{S_i-T_i+D_i}{D_i} \]
を \(998244353\) で割った余りとなるため、これを計算すれば良いです。
\(S_i\) および \(T_i\) は深さ優先探索によって全体で\(O(N)\) で求めることができます。
また、\(\binom{S_i-T_i+D_i}{D_i}=\frac{1}{D_i!}\{(S_i-T_i+D_i)\times(S_i-T_i+D_i-1)\times\cdots\times(S_i-T_i+1)\}\) の値は\(P=998244353\) としてそれぞれ \(O(D_i\log P)\) 、あるいは階乗 \(k!\) を\(1\leq k\leq \max(D_i)\) の範囲で前計算しておくことで \(O(D_i+\log P)\) で求めることができます。よって、全体で時間計算量は \(\displaystyle O\biggl((\sum_{i=1}^N D_i)\log P\biggr)\)、あるいは \(\displaystyle O\biggl((\sum_{i=1}^N D_i)+N\log P\biggr)\) であり、いずれも問題の制約下で十分高速です。よって、この問題を解くことができました。
C++ による実装例:
#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
#define ll long long
#define rep(i, n) for(int i = 0; i < n; ++i)
#define N (int)2e+5
vector<int>e[N];
ll c[N];
ll d[N];
mint ans;
ll dfs(int k){
int sz=e[k].size();
ll s=c[k];
for(int i=0;i<sz;i++)s+=dfs(e[k][i]);
if(s<d[k])ans=0;
else{ //binom (s,d[k])
rep(i,d[k]){
ans*=mint(s-i);
ans*=mint(i+1).inv();
}
}
return (s-d[k]);
}
int main(void){
int n,x;
cin>>n;
for(int i=1;i<n;i++){
cin>>x;
e[x-1].push_back(i);
}
rep(i,n)cin>>c[i];
rep(i,n)cin>>d[i];
ans=1;
dfs(0);
cout<<ans.val()<<endl;
return 0;
}
posted:
last update:
