[洛谷] P3914 染色计数

Description

有一颗(N)个节点的树，节点用(1,2,cdots,N)编号。你要给它染色，使得相邻节点的颜色不同。有(M)种颜色，用(1,2,cdots,M)编号。每个节点可以染(M)种颜色中的若干种，求不同染色方案的数量除以(10^9 + 7)的余数。

Solution

设状态(dp[i][j])为节点(i)，颜色为(j)的方案数，则有

[dp[i][j]=sum_{k=1,k!=j}^ndp[son[i]][k](初始化dp[i][col[i]]=1) ]

然后这个可以(dfs)的时候顺便处理(dp[i][1sim{m}])的和，就可以做到(O(nm))了。

类似的题有P4084 [USACO17DEC]Barn Painting G

Code

#include <bits/stdc++.h>

using namespace std;

const int mod = 1e9 + 7;

int n, m, tot, res, vis[5005], hd[5005], to[10005], nxt[10005], sum[5005], dp[5005][5005];

int read()
{
	int x = 0, fl = 1; char ch = getchar();
	while (ch < '0' || ch > '9') { if (ch == '-') fl = -1; ch = getchar();}
	while (ch >= '0' && ch <= '9') {x = (x << 1) + (x << 3) + ch - '0'; ch = getchar();}
	return x * fl;
}

void add(int x, int y)
{
	tot ++ ;
	to[tot] = y;
	nxt[tot] = hd[x];
	hd[x] = tot;
	return;
}

void dfs(int x, int fa)
{
	for (int i = hd[x]; i; i = nxt[i])
	{
		int y = to[i];
		if (y == fa) continue;
		dfs(y, x);
		for (int p = 1; p <= m; p ++ )
			dp[x][p] = (1ll * dp[x][p] * (sum[y] - dp[y][p]) % mod + mod) % mod;
	}
	for (int p = 1; p <= m; p ++ )
		sum[x] = (sum[x] + dp[x][p]) % mod;
	return;
}

int main()
{
	n = read(); m = read();
	for (int i = 1; i <= n; i ++ )
	{
		int t = read();
		while (t -- )
		{
			int col = read();
			dp[i][col] = 1;
		}
	}
	for (int i = 1; i <= n - 1; i ++ )
	{
		int x = read(), y = read();
		add(x, y); add(y, x);
	}
	dfs(1, 0);
	for (int p = 1; p <= m; p ++ )
		res = (res + dp[1][p]) % mod;
	printf("%d
", res);
	return 0;
}