【树上统计，组合数学】Select from Subtrees

题意概括

给定一棵 $N$ 个节点的有根树（根为 1）。每个节点 $i$ 初始有 $C_i$ 个不同的糖果。 有 $N$ 只松鼠，第 $i$ 只松鼠需要从以 $i$ 为根的子树中挑选 $D_i$ 个糖果。不同松鼠不能选同一个糖果，不同松鼠选到相同的糖果组合视为不同的分配方案。 求所有满足要求的分配方案数，对 998244353 取模。 输入包含树的边、每个节点的 $C_i$ 和 $D_i$。数据范围：$N \le 2 \times 10^5$，$C_i \le 10^9$，$\sum D_i \le 10^6$。 保留样例 1 进行说明：5 个节点，结构为 1 连 2 和 3，3 连 4 和 5。各节点的 $C$ 为 1, 2, 1, 2, 3，$D$ 为 1, 1, 3, 1, 1。输出为 144。

样例分析

在样例 1 中，节点 2、4、5 是叶子节点，它们上面的松鼠只能在自身节点拿糖果。 节点 2 有 2 个糖果，需拿 1 个，方案数为 $\binom{2}{1} = 2$。 节点 4 有 2 个糖果，需拿 1 个，方案数为 $\binom{2}{1} = 2$。 节点 5 有 3 个糖果，需拿 1 个，方案数为 $\binom{3}{1} = 3$。 节点 3 的子树包含节点 3、4、5，共有 $1+2+3=6$ 个糖果。节点 4 和 5 的松鼠各自拿走了 1 个，还剩下 $6-2=4$ 个糖果。松鼠 3 需要拿 3 个，方案数为 $\binom{4}{3} = 4$。 节点 1 的子树包含整棵树，共有 9 个糖果。其他松鼠共计拿走 6 个，留给松鼠 1 的剩余糖果为 $9-6=3$ 个。松鼠 1 需要拿 1 个，方案数为 $\binom{3}{1} = 3$。 根据乘法原理，总方案数为 $2 \times 2 \times 3 \times 4 \times 3 = 144$。

题目分析

题目要求各个松鼠在自己的子树内选糖果，且糖果各不相同。由于子树内的糖果对于子树根节点及以上的祖先节点是完全等价的，可以采用自底向上的方式计算方案数。 优先满足深度较大的节点需求。对于节点 $u$ 的松鼠，它需要从 $u$ 的子树中选出 $D_u$ 个糖果。当处理到 $u$ 时，它所有的真后代（子树中除 $u$ 外的节点）已经从这片区域中预定了它们需要的糖果。 设 $u$ 的子树中初始总糖果数为 $W_u$，而 $u$ 的所有真后代总共需要拿走 $R_u$ 个糖果。真后代的选取限制在它们各自的子树内，只要数量足够，它们挑走特定的 $R_u$ 个糖果后，剩下的 $W_u - R_u$ 个糖果对于松鼠 $u$ 而言没有任何限制，可以任选。 对每个节点 $u$，当前可用的糖果总数即为 $W_u - R_u$。松鼠 $u$ 的局部方案数就是 $\binom{W_u - R_u}{D_u}$。遍历整棵树时，如果出现子树总糖果数小于总需求数（即 $W_u < R_u + D_u$），直接输出 0 并结束程序；否则将所有节点的局部方案数相乘即可。 组合数计算需要特别处理。$W_u$ 的累加值可能达到 $2 \times 10^{14}$，而取模的质数 $p = 998244353$ 大于 $10^6$。根据卢卡斯定理，$\binom{n}{m} \equiv \binom{n \bmod p}{m \bmod p} \times \binom{n/p}{m/p} \pmod p$。因为 $D_u \le 10^6 < p$，所以 $D_u / p = 0$，右侧第二项 $\binom{n/p}{0} = 1$。公式直接退化为计算 $\binom{n \bmod p}{D_u} \pmod p$。将可用糖果数量对 $p$ 取模，若取模后数值小于 $D_u$ 则返回 0，否则在 $O(D_u)$ 时间内通过展开分子分母计算组合数。

参考实现

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

复杂度分析

时间复杂度：树的后序遍历需要 $O(N)$ 时间。节点 $u$ 的组合数计算需要 $O(D_u)$ 时间，整棵树累计计算次数为 $\sum D_i \le 10^6$。整体时间复杂度为 $O(N + \sum D_i)$。 空间复杂度：使用邻接表存储树的边，维护每个节点的糖果信息，空间复杂度为 $O(N)$。

https://atcoder.jp/contests/abc459/tasks/abc459_e