题面

有一颗 $n$ 个点的树，第 $i$ 个点有点权 $a_i$ 。

现在要断掉一些边，记一个连通块的权值为连通块中所有点的点权和，一个断边方案的权值为所有连通块的权值的乘积的平方。

求所有不同的断边方案的权值和对 $10^9+7$ 取模后的结果。

$n\leq5\times10^5$

题解

直接组合意义：每个连通块 $\sum a_i$ 个球，对每个连通块选两个球（放回）染色的方案数。

树上 DP 即可：设 $f_{u,i}$ 表示 $u$ 的子树中选 $i$ 个球的方案数。转移：

1. 不删边： $f_{u,i}\leftarrow_{v\in\text{son}(u),j\leq i} f_{v,j}\times a_u^{i-j}\times\binom{i}{j}$ ；
2. 删边： $f_{u,i}\leftarrow_{v\in\text{son}(u)}f_{v,2}\times a_u^i$ 。

$\mathcal{O}(n)$ 。

另外如果是 $k$ 次方可以用多项式做法 $\mathcal{O}(nk\log k)$ 。

代码

const int N = 5e5 + 5;
const int mod = 1e9 + 7;

inline ll Read() {
	ll x = 0, f = 1;
	char c = getchar();
	while (c != '-' && (c < '0' || c > '9')) c = getchar();
	if (c == '-') f = -f, c = getchar();
	while (c >= '0' && c <= '9') x = (x << 3) + (x << 1) + c - '0', c = getchar();
	return x * f;
}

namespace Main {
	int n;
	ll a[N];
	ll f[N][5], C[5][5];
	vector <int> e[N];
	void dfs (int u, int fa) {
		ll g[3] = {0};
		for (int v: e[u]) {
			if (v == fa) continue;
			g[0] = f[u][0], g[1] = f[u][1], g[2] = f[u][2];
			dfs(v, u);
			f[u][0] = f[u][1] = f[u][2] = 0;
			for (int i = 0; i <= 2; i++) {
				for (int j = 0; j <= i; j++)
					f[u][i] = (f[u][i] + 1ll * f[v][j] * g[i - j] % mod * C[i][j] % mod) % mod;
				f[u][i] = (f[u][i] + 1ll * f[v][2] * g[i] % mod) % mod;
			}
		}
	}
	int main () {
		n = Read();
		for (int i = 1; i <= n; i++) a[i] = Read(), f[i][0] = 1, f[i][1] = a[i], f[i][2] = a[i] * a[i] % mod;
		for (int i = 0; i <= 3; i++) C[i][0] = 1, C[i][i] = 1;
		for (int i = 1; i <= 3; i++)
			for (int j = 1; j < i; j++) C[i][j] = C[i - 1][j] + C[i - 1][j - 1];
		for (int i = 1; i < n; i++) {
			int u = Read(), v = Read();
			e[u].emplace_back(v);
			e[v].emplace_back(u);
		}
		dfs (1, 0);
		printf ("%lld\n", f[1][2]);
		return 0;
	}
}

int main () {
	freopen("hard.in", "r", stdin);
	freopen("hard.out", "w", stdout);
	Main::main();
	return 0;
}