杭电多校第三场 1005 Little W and Contest 题解

题目大意

总共有n个人，每个人都有自己的权值1或者2。

组成一个队伍至少需要三个人，最少需要两个2。每个队伍不能存在任意两个互相认识的人。

输入n-1组关系，使两个人熟悉，保证之前是都不认识的。输出每一次熟悉之后剩下的能组成队伍的个数。

解题思路

这个题目其实看懂了不难。

首先全部个数就是$C_{cnt2}^{3}+C_{cnt1}^1*C_{cnt2}^2$。

（在求这个的时候我又用了取模。。我又忘了除法不能同余。我自裁。。

之后就只需要在这个基础上不断减少就可以了。

用c1和c2两个数组来维持并查集中不同集合中的为1和为2的个数。

因为之前都是互不认识的，所以必定根节点不同。

（之前就没注意到，想了好久要不要判断是否根节点不同。。

之后的就简单了，考虑到四种情况就可以了。

1、1 2 2

2、2 1 2

3、2 2 1

4、2 2 2

把这四种情况减去就可以了。

完整代码

#include <bits/stdc++.h>
#include <iostream>
#include <cstdio>
#include <queue>
#include <cstring>
#include <cmath>
#include <stack>
#include <map>
#include <string>
#include <vector>
#include <algorithm>
#include <sstream>
#include <unordered_map>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;

const int N = 1e5 + 5, mod = 1e9 + 7;
int fa[N];
ll c1[N], c2[N];
ll cnt1, cnt2;
ll ans;

int find(int son) {
	int root = son;
	while (fa[root] != root)	root = fa[root];
	while (son != root) {
		int t = fa[son];
		fa[son] = root;
		son = t;
	}
	return root;
}

void join(int a,int b) {
	int rt1 = find(a), rt2 = find(b);
	ans -= c2[rt1] * c2[rt2] * (cnt2 - c2[rt1] - c2[rt2]);
	ans -= c2[rt1] * c2[rt2] * (cnt1 - c1[rt1] - c1[rt2]);
	ans -= c1[rt1] * c2[rt2] * (cnt2 - c2[rt1] - c2[rt2]);
	ans -= c2[rt1] * c1[rt2] * (cnt2 - c2[rt1] - c2[rt2]);
	fa[rt1] = rt2;
	c1[rt2] += c1[rt1], c2[rt2] += c2[rt1];
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);cout.tie(0);
	int t; cin >> t;
	while (t--) {
		int n; cin >> n;
		cnt1 = cnt2 = 0;
		for (int i = 1; i <= n; i++) {
			int x; cin >> x;
			fa[i] = i;
			if (x == 1)	c1[i] = 1, c2[i] = 0, cnt1++;
			else c1[i] = 0, c2[i] = 1, cnt2++;
		}
		ans = cnt2 * (cnt2 - 1) * (cnt2 - 2) / 6 + cnt2 * (cnt2 - 1) * cnt1 / 2;
		cout << ans % mod << endl;
		for (int i = 1; i < n; i++) {
			int u, v; cin >> u >> v;
			join(u, v);
			cout << ans % mod << endl;
		}
	}
	return 0;
}