鸡玩炸弹人

原题链接

TAG：思维，并查集

思路

对于任意一个连通块 \(i\)，它的大小为 \(siz_i\)，无论炸弹如何放置，都有 \(siz_i\times siz_i\) 个起点与终点的选择。

证明

考虑联通块是一颗树的情况，可以先从S不放炸蛋的走到T，然后从T出发，按照类似dfs的方式遍历这棵树，在回溯时选择放炸蛋即可做到放完所有炸蛋最终回到T。

因此记录第 \(i\) 个连通块的大小 \(siz_i\)，有炸弹的连通块的数量为 \(comp\)。则：

• 当 \(comp=0\)，输出 \(\sum siz_i^2\)；
• 当 \(comp=1\)，输出有炸弹的连通块 j 的 \(siz_j^2\)；
• 当 \(comp\ge 2\)，输出 \(0\)。

代码

#include <bits/stdc++.h>
using namespace std;
using LL = long long;

struct DSU {
    std::vector<int> f, siz;
    DSU(int n) : f(n), siz(n, 1) {
        std::iota(f.begin(), f.end(), 0);
    }
    int leader(int x) {
        while (x != f[x]) x = f[x] = f[f[x]];
        return x;
    }
    bool same(int x, int y) {
        return leader(x) == leader(y);
    }
    bool merge(int x, int y) {
        x = leader(x);
        y = leader(y);
        if (x == y) return false;
        siz[x] += siz[y];
        f[y] = x;
        return true;
    }
    int size(int x) {
        return siz[leader(x)];
    }
};

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    int n, m;
    cin >> n >> m;

    DSU dsu(n);
    while (m--) {
        int u, v;
        cin >> u >> v;
        u--;
        v--;
        dsu.merge(u, v);
    }

    int comp = 0; // 记录有炸弹的连通块数量
    vector<int> s(n);
    for (int i = 0; i < n; i++) {
        int c;
        cin >> c;
        if (c) {
            int x = dsu.leader(i);
            if (!s[x]) comp++;
            s[x] += c;
        }
    }

    LL res = 0;
    for (int i = 0; i < n; i++) {
        if (dsu.leader(i) == i) { // 找到连通块
            if (comp - (s[i] > 0) == 0) {
                int s = dsu.size(i);
                res += 1ll * s * s;
            }
        }
    }
    cout << res << "\n";

    return 0;
}