题目要求各个松鼠在自己的子树内选糖果,且糖果各不相同。由于子树内的糖果对于子树根节点及以上的祖先节点是完全等价的,可以采用自底向上的方式计算方案数。
优先满足深度较大的节点需求。对于节点 u 的松鼠,它需要从 u 的子树中选出 Du 个糖果。当处理到 u 时,它所有的真后代(子树中除 u 外的节点)已经从这片区域中预定了它们需要的糖果。
设 u 的子树中初始总糖果数为 Wu,而 u 的所有真后代总共需要拿走 Ru 个糖果。真后代的选取限制在它们各自的子树内,只要数量足够,它们挑走特定的 Ru 个糖果后,剩下的 Wu−Ru 个糖果对于松鼠 u 而言没有任何限制,可以任选。
对每个节点 u,当前可用的糖果总数即为 Wu−Ru。松鼠 u 的局部方案数就是 (DuWu−Ru)。遍历整棵树时,如果出现子树总糖果数小于总需求数(即 Wu<Ru+Du),直接输出 0 并结束程序;否则将所有节点的局部方案数相乘即可。
组合数计算需要特别处理。Wu 的累加值可能达到 2×1014,而取模的质数 p=998244353 大于 106。根据卢卡斯定理,$\binom{n}{m} \equiv \binom{n \bmod p}{m \bmod p} \times \binom{n/p}{m/p} \pmod p$。因为 Du≤106<p,所以 Du/p=0,右侧第二项 (0n/p)=1。公式直接退化为计算 (Dunmodp)(modp)。将可用糖果数量对 p 取模,若取模后数值小于 Du 则返回 0,否则在 O(Du) 时间内通过展开分子分母计算组合数。
参考实现
#include<bits/stdc++.h>
using namespace std;
using ll = long long;
const int MOD = 998244353;
const int MX = 200005;
vector<int> adj[MX];
ll c[MX], d[MX], sc[MX], sd[MX];
ll ans = 1;
ll qpw(ll a, ll b) {
ll r = 1;
while (b) {
if (b & 1) r = r * a % MOD;
a = a * a % MOD;
b >>= 1;
}
return r;
}
ll inv(ll a) {
return qpw(a, MOD - 2);
}
ll get(ll n, ll k) {
n %= MOD;
if (n < k) return 0;
ll num = 1, den = 1;
for (ll i = 0; i < k; i++) {
num = num * (n - i) % MOD;
den = den * (i + 1) % MOD;
}
return num * inv(den) % MOD;
}
void dfs(int u) {
sc[u] = c[u];
sd[u] = d[u];
for (int v : adj[u]) {
dfs(v);
sc[u] += sc[v];
sd[u] += sd[v];
}
if (sc[u] < sd[u]) {
cout << 0 << "\n";
exit(0);
}
ll rem = sc[u] - sd[u] + d[u];
ans = ans * get(rem, d[u]) % MOD;
}
int main() {
ios::sync_with_stdio(0);
cin.tie(0);
cout.tie(0);
int n;
cin >> n;
for (int i = 2; i <= n; i++) {
int p;
cin >> p;
adj[p].push_back(i);
}
for (int i = 1; i <= n; i++) cin >> c[i];
for (int i = 1; i <= n; i++) cin >> d[i];
dfs(1);
cout << ans << "\n";
return 0;
}