Wannafly挑战赛19C：多彩的树

https://ac.nowcoder.com/acm/contest/5556/A

代码今天晚上补完

 

#include<iostream>
#include<cstring>
#include<algorithm>
#include<queue>
#include<vector>
using namespace std;
typedef long long ll;
ll mod = 1e9 + 7;

const int maxn = 1e5 + 11;

int ans[maxn];
int list[maxn];//颜色
int num[200000];
vector<int>G[maxn];
void add(int x, int y) {
	G[x].push_back(y);
}
int vis[maxn];
ll siz[maxn];
ll cnt[maxn];


int dfs(int x, int clor) {//clor是一个状态
	vis[x] = 1;
	siz[x] = 1;
	for (int i = 0; i < G[x].size(); i++) {
		int p = G[x][i];
		if (vis[p] == 0 && (clor & (1 << list[p]))) {
			dfs(p, clor);
			cnt[clor] += siz[x] * siz[p];
			siz[x] += siz[p];
		}
	}
	cnt[clor]++;
	return 0;
}

ll k_q(ll a, ll b) {
	ll res = 1;
	while (b) {
		if (b & 1) {
			res = (a*res) % mod;
		}
		b >>= 1;
		a = (a*a) % mod;
	}
	return res;
}

int main() {
	for (int i = 1; i < (1<<10) ; i++) {
		for (int j = 0; j < 10; j++) {
			if (i & (1 << j)) num[i]++;
		}
	}

	int n, k;
	scanf("%d %d", &n, &k);
	for (int i = 1; i <= n  ; i++){
		scanf("%d", &list[i]);
		list[i]--;
	}

	int be, en;
	for (int i = 1; i < n; i++) {
		scanf("%d%d", &be, &en);
		add(be, en);
		add(en, be);
	}
	for (int i = 1; i < (1<<k); i++) {
		for (int j = 0; j <= n; j++) vis[j] = 0;
		for (int j = 1; j <= n; j++) {
			if (vis[j] == 0 && (i&(1<<list[j]))) {
				dfs(j, i);
			}
		}
	} 

	for (int i = 1; i < (1 << k); i++) {
		for (int j = 1; j < (1 << k); j++) {
			if (i == j) continue;
			if ((i&j) == j) {
				cnt[i] -= cnt[j];
			}
		}
		ans[num[i]] += cnt[i];
	}

	ll a = 0;
	for (int i = 1; i <= k; i++) {
		a = (a + ans[i] * k_q(131, i) % mod) % mod;
	}
	printf("%lld\n", a);

	return 0;
}