2021 陕西省赛 D - Disease // 树形概率dp

题目来源：2021年ICPC国际大学生程序设计竞赛暨陕西省第九届大学生程序设计竞赛

题目链接：https://ac.nowcoder.com/acm/contest/35232/D

题意

给定一棵根为 (1)，大小为 (n) 的树：树上每个结点 (i) 初始有 (frac{large p_i}{large q_i}) 的概率从外界染病，对于每条边 ({ u, v, a_j,b_j })，表示 (u) 和 (v) 相连，且当其中一个点染病时，另一个点被传染的概率为 (frac{large a_j}{large b_j}) .

定义灾难等级为：染病结点的最小深度（根节点 (1) 的深度为 (1)）。

求：灾难等级的期望，答案对 (10^9 + 7) 取模。

数据范围：(1 le n le 10^5)，(0 le p_i,q_i,a_j,b_j le 10^6).

思路：树形概率dp

考虑枚举灾难等级，即存在染病结点的最小深度 (i)，若其概率为 (g_i)，那么其对答案的贡献为 (i times g_i) .

于是我们需要计算出 (g_i) ：

(g_i) 的事件实际上是，前 (i-1) 层不存在染病结点，且第 (i) 层存在染病结点。

更进一步分解为，前 (i-1) 层的结点初始均不染病（记为(P_1)），第 (i) 层存在自己染病但未传染给父亲的结点（记为(P_2)），即 (g_i = P_1 times P_2) .

容易得到，(P_1 = prodlimits_{depth_{large j} < i} (1-frac{large p_j}{large q_j})) .

对于(P_2) 的事件，其等价于：第 (i) 层的结点均不是自己染病并传染给父亲（即每个结点要么不染病，要么染病但不传染给父亲），并且，并非所有结点都不染病。则有，(P_2 = prodlimits_{depth_{large j} = i} (1 - f_j times frac{large a_j}{large b_j}) - prodlimits_{depth_{large j} = i}f_j) ，其中：(f_j) 为 每个结点 (j) 最终不染病的概率，(frac{large a_j}{large b_j}) 为 结点 (j) 传染给父亲的概率。

于是我们还需要先计算出 (f_j) ：

在我们计算 (g_i) 时，实际上限定了前 (i-1) 层的结点均初始不染病，因此我们是不需要考虑往下的传染的，只需要考虑往上的传染。

那么，(f_j) 的事件就是，结点 (j) 初始不染病，并且每个儿子都是要么不染病要么染病且不传染。

于是有，(f_j = (1 - frac{large p_j}{large q_j}) times prodlimits_{k in son_{large j}}(f_k + (1-f_k)·(1-frac{large a_k}{large b_k})))，其中：(frac{large a_k}{large b_k}) 为 结点 (k) 传染给其父亲 (j) 的概率。

总的来说，需要先做一个树形dp求出 (f_j)，再按层数由小到大求出 (g_i)，并统计其对答案的贡献 (i times g_i)，最终答案为：

(ans = sumlimits_{i=1}^{maxDepth}i times prodlimits_{depth_{large j} < i} (1-frac{large p_j}{large q_j}) times (prodlimits_{depth_{large j} = i} (1 - f_j · frac{large a_j}{large b_j}) - prodlimits_{depth_{large j} = i}f_j)) .

时间复杂度：(O(nlogP))，(P = 10^9 + 7).

代码

#include <bits/stdc++.h>
#define endl 'n'
#define int long long
using namespace std;

const int N = 100010;
const int mod = 1e9 + 7;

int n, dep[N], f[N], p[N], q[N];
vector<pair<int,int>> g[N];
vector<int> vec[N];

int qmi(int m, int k, int p) 
{
    int res = 1;
    while(k) {
        if(k & 1) res = res * m % p;
        m = m * m % p;
        k >>= 1;
    }
    return res;
}

int getmod(int x) 
{
    return (x % mod + mod) % mod;
}

int dfs(int u, int fa) {
    if(~f[u]) return f[u];
    vec[dep[u]].push_back(u);
    f[u] = getmod(1 - p[u]);
    for(auto e : g[u]) {
        int v = e.first, w = e.second;
        if(v == fa) continue;
        dep[v] = dep[u] + 1, q[v] = w;
        int p1 = dfs(v, u);
        int p2 = getmod(1 - dfs(v, u)) * getmod(1 - w) % mod;
        f[u] = f[u] * (p1 + p2) % mod;
    }
    return f[u];
}

signed main()
{
    cin.tie(0);
    ios::sync_with_stdio(false);

    cin >> n;
    for(int i = 1; i <= n; ++ i) {
        int x, y;
        cin >> x >> y;
        p[i] = x * qmi(y, mod - 2, mod) % mod;
    }
    for(int i = 1; i < n; ++ i) {
        int u, v, a, b;
        cin >> u >> v >> a >> b;
        int w = a * qmi(b, mod - 2, mod) % mod;
        g[u].push_back({v, w}), g[v].push_back({u, w});
    }

    memset(f, -1, sizeof f);
    dep[1] = 1, dfs(1, -1);

    int ans = 0, P = 1;
    for(int i = 1; i <= n; ++ i) {
        if(!vec[i].size()) break;
        int res = P, P1 = 1, P2 = 1;
        for(auto x : vec[i]) {
            P = P * getmod(1 - p[x]) % mod;
            P1 = P1 * getmod(1 - getmod(1 - f[x]) * q[x] % mod) % mod;
            P2 = P2 * f[x] % mod;
        }
        res = res * getmod(P1 - P2) % mod;
        ans = (ans + res * i) % mod;
    }
    cout << ans << endl;

    return 0;
}