C. Helga Hufflepuff‘s Cup

题意

给一棵树， 并且给m个颜色， 让你染色， 有一种特殊颜色k，特殊颜色最短染x个点， 问有多少种染色情况。 结果取模。

分析

f[ i ] [ j ] [ k] k 可以取0， 1， 2， 取0表示 当前以i为根的树用了j个特殊颜色，并且当前是特殊颜色的情况。 1表示 当前染色小于k，的所有情况， 2表示当前染色大于k的所有情况。
需要新开一个数组，否则当前的情况会对后面的情况产生影响。

代码

#include<bits/stdc++.h>
using namespace std;
typedef long long LL;
const LL N = 1e5+5, M = N * 2, mod = 1e9+7;
 
int n, m;
int x, k;
 
LL h[N], e[M], ne[M], w[M], idx;
LL f[N][12][5], sz[N], f1[12][5];
// i为根 选j个最高点  0选k  1选k小 2选比k大
 
void add(int a, int b)
{
	e[idx] = b, ne[idx] = h[a], h[a] = idx ++;
}
 
void dfs(int u, int fa)
{
	sz[u] = 1;
	f[u][1][0] = 1; // 当前放特殊色， 有一种
	f[u][0][1] = k-1; // 当前染特殊色 种类
	f[u][0][2] = m-k;// 当前放比k大的情况
	
	for(int i = h[u]; ~i; i = ne[i])
	{
		int j = e[i];
		if(j == fa) continue;
		
		dfs(j, u);
		sz[u] += sz[j];
		
		memset(f1, 0, sizeof f1); // 将当前为根的数组还原
		for(int jj = 0; jj <= sz[u] && jj <= x; jj ++)
			for(int kk = 0; kk <= sz[j]&&jj+kk <= x; kk ++)
			{
				f1[jj+kk][0] = (f1[jj+kk][0] + f[u][jj][0]*f[j][kk][1]%mod)%mod;
				// 当前为根 特殊点选了kk+jj个当前选的也是特殊点
				f1[jj+kk][1] = (f1[jj+kk][1] + f[u][jj][1]*(f[j][kk][0]+f[j][kk][1]+f[j][kk][2])%mod)%mod;
				
				f1[jj+kk][2] = (f1[jj+kk][2] + f[u][jj][2]*(f[j][kk][1]+f[j][kk][2])%mod)%mod;
			}
		
		for(int jj = 0; jj <= sz[u]&& jj <= x; jj ++)
			for(int kk = 0; kk < 3; kk ++)
				f[u][jj][kk] = f1[jj][kk];
	}
 
}
 
int main()
{
	memset(h, -1, sizeof h);
	cin >> n >> m; // m 种颜色
	for(int i = 1; i < n; i ++)
	{
		int a, b;
		cin >> a >> b;
		add(a, b), add(b, a);
	}
	cin >> k >> x; // k为特殊点， 数量x
	
	dfs(1, -1);
	
	LL ans = 0;
	for(int i = 0; i <= x; i ++)
	{
		ans = (ans + f[1][i][0])%mod;
		ans = (ans + f[1][i][1])%mod;
		ans = (ans + f[1][i][2])%mod;
	}
	cout << ans << endl;
	return 0;
}