E - Select from Subtrees 解説 by en_translator
The order the squirrels choose candies do not affect the answer, so we will decide candies picked by squirrels in post-order of the vertices in the tree \(T\).
Let \(S_i\) and \(T_i\) be the sum of \(C_j\) and \(D_j\), respectively, over vertices \(j\) in the subtree rooted at vertex \(i\). In the post-order of the vertices, any descendants of vertex \(i\) comes earlier than vertex \(i\), and ancestors later, so there are \(S_i-(T_i-D_i)\) candidate candies (no pun intended) that squirrel \(i\) may choose, regardless of the prior squirrels’ choice, from which squirrel \(i\) chooses \(D_i\) of them. Thus, there are \(\binom{S_i-T_i+D_i}{D_i}\) ways for squirrel \(i\) to choose candies, independent of earlier squirrels’ decisions. Here, if \(S_i-T_i+D_i<D_i\), or \(S_i<T_i\), then squirrel \(i\) cannot choose candies as demanded, so the answer is \(0\).
The answer is the product of this value over \(i=1,2,\ldots,N\), that is:
\[ \prod_{i=1}^n \binom{S_i-T_i+D_i}{D_i}, \]
divided by \(998244353\). All that left is to calculate this.
\(S_i\) and \(T_i\) can be found in a total of \(O(N)\) time with Depth-First Search (DFS).
Also, the value of \(\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)\}\) can be found in \(O(D_i\log P)\) time, or in \(O(D_i+\log P)\) time by precomputing the factorials \(k!\) for all \(1\leq k\leq \max(D_i)\), where \(P=998244353\). Hence, the total time complexity is \(\displaystyle O\biggl((\sum_{i=1}^N D_i)\log P\biggr)\) or \(\displaystyle O\biggl((\sum_{i=1}^N D_i)+N\log P\biggr)\), which is both fast enough under the constraints of the problem. Thus, the problem has been solved.
Sample code in 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;
}
投稿日時:
最終更新: