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_j, v_j, a_j,b_j \}\)，表示 \(u_j\) 和 \(v_j\) 相连，且当其中一个点染病时，另一个点被传染的概率为 \(\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 = \prod\limits_{depth_{\large j} < i} (1-\frac{\large p_j}{\large q_j})\) .

对于\(P_2\) 的事件，其等价于：第 \(i\) 层的结点均不是自己染病并传染给父亲（即每个结点要么不染病，要么染病但不传染给父亲），并且，并非所有结点都不染病。则有，\(P_2 = \prod\limits_{depth_{\large j} = i} (1 - f_j \times \frac{\large a_j}{\large b_j})\ - \prod\limits_{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 \prod\limits_{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\ = \sum\limits_{i=1}^{maxDepth}i\ \times \prod\limits_{depth_{\large j} < i} (1-\frac{\large p_j}{\large q_j}) \times (\prod\limits_{depth_{\large j} = i} (1 - f_j · \frac{\large a_j}{\large b_j})\ - \prod\limits_{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;
}