連結

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

題意

人分兩類，1類2類
隊有三人，至少兩個2類
共n人，一開始互不認識
人有一特性，朋友的朋友也是朋友
隊有一特性，朋友不能共處
輸入兩個人n-1次，表示將此兩個人介紹為朋友，請在每次介紹前輸出有多少種組隊方式
最後一行輸出介紹完有多少種組隊方式

思路

經典題目朋友的朋友也是朋友，並查集維護連通塊，連通塊內的元素不能組隊
容易發現每次將兩個連通塊合併後答案會減少，所以計算兩個連通塊的負貢獻即可
假設A，B為要合併的兩連通塊，C為其他連通塊，則所有對答案有影響的取法為

將以上四種情況的計數減去即可，具體見程式碼

程式碼

#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;
}