公式

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;
}

投稿日時:
最終更新: