NC19782 Tree

题目链接

题目

题目描述

修修去年种下了一棵树，现在它已经有n个结点了。
修修非常擅长数数，他很快就数出了包含每个点的连通点集的数量。
澜澜也想知道答案，但他不会数数，于是他把问题交给了你。

输入描述

第一行一个整数 $n (1≤ n ≤ 10^6)$ ，接下来n-1行每行两个整数 $a_i,b_i$ 表示一条边 $(1≤ a_i,b_i≤ n)$ 。

输出描述

输出n行，每行一个非负整数。第i行表示包含第i个点的连通点集的数量对 $10^9+7$ 取模的结果。

示例1

输入

6
1 2
1 3
2 4
4 5
4 6

输出

12
15
7
16
9
9

题解

知识点：树形dp，计数dp。

一道树形dp，注意到要求每个点的贡献，因此需要二次扫描+换根。

第一次先随便选个点，这里考虑以 $1$ 为根开始dp。设 $dp[u]$ 表示 $u$ 为根的子树中过 $u$ 的连通点集数量，转移方程为：

$$dp[u] = \prod (dp[v_i] + 1)$$

表示合并以子节点 $v$ 为根的子树的方案数，加一是不选 $v$ 这个子树的一种可能，乘法原理合并。

第二次，在得到了各个子树答案后，可以顺推完整答案。设 $f[u]$ 表示过 $u$ 的连通点集数量，转移方程为：

$$ans[v] = dp[v] \cdot \bigg( \frac{ans[u]}{dp[v]+1}+1 \bigg)$$

表示从父节点的答案 $ans[u]$ 删去选或不选 $v$ 这一支路选的答案 $dp[v]+1$ ，再加上不选父节点 $u$ 的答案一种，与 $v$ 支路的答案进行乘法原理合并，最终得到 $v$ 的完整答案 $ans[v]$ ，其中需要求 $dp[v]+1$ 的逆元。

但有个大坑，$dp[v]+1 \equiv 0 \quad (\mod 10^9+7)$ 时是没有逆元的，这时候只能重新求一遍 $\frac{ans[u]}{dp[v]+1}$ 。

如果重新开始求，把 $u$ 作为根节点再求一次 $dp[u]$ ，这样复杂可能会退化到枚举根的程度。

因此，考虑设 $f[v] = \frac{ans[u]}{dp[v]+1}$ ，于是 $ans[v] = dp[v] \cdot (f[v]+1)$ ，$f[v]$ 可以在每次求 $ans[v]$ 前先求出。如果 $dp[v]+1$ 有逆元 $(dp[v]+1)^{-1}$ ，则 $f[v] = ans[u] \cdot (dp[v]+1)^{-1}$ ；如果没有，则 $f[v] = \frac{ans[u]}{dp[v]+1} = \frac{dp[u] }{dp[v]+1} \cdot (f[u]+1)$ ，其中 $\dfrac{dp[u] }{dp[v]+1}$ 只需要求一遍 $\prod (dp[v_i]+1),v_i \not= v$ 。注意初始时 $f[1] = 0$ ，因为此时 $ans[1] = dp[1]$ 。

时间复杂度 $O(n)$

空间复杂度 $O(n)$

代码

#include <bits/stdc++.h>
#define ll long long
 
using namespace std;
 
const int mod = 1e9 + 7;
 
struct Graph {
    int n;
    struct edge {
        int to, nxt;
    };
    vector<edge> e;
    vector<int> h;
 
    explicit Graph(int _n) :n(_n), h(_n + 1, -1) {}
 
    void add(int u, int v) {///加边
        e.push_back(edge{ v,h[u] });///边结束节点，边出发节点的上一条边在e中下标，边权
        h[u] = e.size() - 1;///上一条边的下标
    }
};
Graph g(1000007);
ll dp[1000007], f[1000007];
 
ll qpow(ll a, ll k) {
    ll ans = 1;
    while (k) {
        if (k & 1) ans = ans * a % mod;
        k >>= 1;
        a = a * a % mod;
    }
    return ans;
}///求逆元，但得到0时代表没有逆元，因此要重新求
 
void dfs1(int u, int fa) {
    dp[u] = 1;
    for (int i = g.h[u];~i;i = g.e[i].nxt) {
        int v = g.e[i].to;
        if (v == fa) continue;
        dfs1(v, u);
        dp[u] = dp[u] * (dp[v] + 1) % mod;
    }
}///处理出子树中过根节点的答案dp[u]
 
void dfs2(int u, int fa) {
    for (int i = g.h[u];~i;i = g.e[i].nxt) {
        int v = g.e[i].to;
        if (v == fa) continue;
        int inv = qpow(dp[v] + 1, mod - 2);
        if (inv) f[v] = dp[v] * (f[u] * inv % mod + 1) % mod;///inv可以用
        else {
            int tmp = 1;
            for (int j = g.h[u];~j;j = g.e[j].nxt) {
                int vt = g.e[j].to;
                if (vt == fa || vt == v) continue;
                tmp = tmp * (dp[vt] + 1) % mod;
            }
            f[v] = dp[v] * (tmp + 1) % mod;
        }///不能用就重新算一下f[u]/(dp[v]+1)
        dfs2(v, u);
    }
}///通过上一次处理的答案dp[u]，求出完整答案f[v] = dp[v] * (f[u]/(dp[v]+1) + 1)
///表示自己子树的答案与父节点除了自己的答案加上不选父节点的答案相互匹配
 
int main() {
    std::ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    int n;
    cin >> n;
    for (int i = 1;i < n;i++) {
        int u, v;
        cin >> u >> v;
        g.add(u, v);
        g.add(v, u);
    }
    dfs1(1, 0);
    f[1] = dp[1];///初始化
    dfs2(1, 0);
    for (int i = 1;i <= n;i++) cout << f[i] << '\n';
    return 0;
}