2020 Multi-University Training Contest 3 1005- Little W and Contest

链接

http://acm.hdu.edu.cn/showproblem.php?pid=6795

题意

人分两类，1类2类
队有三人，至少两个2类
共n人，一开始互不认识
人有一特性，朋友的朋友也是朋友
队有一特性，朋友不能共处
输入两个人n-1次，表示将此两个人介绍为朋友，请在每次介绍前输出有多少种组队方式
最后一行输出介绍完有多少种组队方式

思路

经典题目朋友的朋友也是朋友，并查集维护连通块，连通块内的元素不能组队
容易发现每次将两个连通块合并后答案会减少，所以计算两个连通块的负贡献即可
假设A，B为要合并的两连通块，C为其他连通块，则所有对答案有影响的取法为

A	B	C
2	1	2
1	2	2
2	2	2
2	2	1

将以上四种情况的计数减去即可，具体见代码

代码

#include <bits/stdc++.h>
#define inf 0x3f3f3f3f
#define ms(a) memset(a, 0, sizeof(a))
#define repu(i, a, b) for (int i = a; i < b; i++)
#define repd(i, a, b) for (int i = a; i > b; i--)
using namespace std;
typedef long long ll;
typedef long double ld;

const int M = int(1e6) + 5;
const int mod = int(1e9) + 7;

ll two;
ll ans;
int n;
int cc[M][2];

int fa[M];
void init() {
    for (int i = 0; i <= n; i++) {
        fa[i] = i;
        cc[i][0] = cc[i][1] = 0;
    }
}
int find(int x) { return fa[x] == x ? x : fa[x] = find(fa[x]); }
int merge(int x, int y) {
    int xx = find(x);
    int yy = find(y);
    fa[xx] = yy;

    ans -= ((ll)cc[xx][0] * cc[yy][1] + (ll)cc[xx][1] * cc[yy][0] +
            (ll)cc[xx][1] * cc[yy][1]) *
           (two - cc[xx][1] - cc[yy][1]);
    ans -= (ll)cc[xx][1] * cc[yy][1] * (n - two - cc[xx][0] - cc[yy][0]);
    cc[yy][0] += cc[xx][0];
    cc[yy][1] += cc[xx][1];

    cout << ans % mod << endl;
}
int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    int t;
    cin >> t;
    while (t--) {
        cin >> n;
        init();

        two = 0;
        repu(i, 1, n + 1) {
            int x;
            cin >> x;
            if (x == 2) {
                two++;
            }
            cc[i][x - 1]++;
        }

        ans = two * (two - 1LL) / 2 * (n - two) +
              two * (two - 1LL) / 2 * (two - 2) / 3;
        cout << ans % mod << endl;

        repu(i, 1, n) {
            int u, v;
            cin >> u >> v;
            merge(u, v);
        }
    }
    return 0;
}