Official

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: