solution.h

/*
Code generated by https://github.com/goodstudyqaq/leetcode-local-tester
*/
#if __has_include("../utils/cpp/help.hpp")
#include "../utils/cpp/help.hpp"
#elif __has_include("../../utils/cpp/help.hpp")
#include "../../utils/cpp/help.hpp"
#else
#define debug(...) 42
#endif

typedef pair<int, int> pii;
struct DSU {
    std::vector<int> f;
    DSU(int n) : f(n) { std::iota(f.begin(), f.end(), 0); }
    int leader(int x) {
        if (f[x] == x) {
            return x;
        }
        int y = leader(f[x]);
        f[x] = y;
        return f[x];
    }
    bool merge(int a, int b) {
        auto x = leader(a);
        auto y = leader(b);
        if (x == y) {
            return false;
        }
        f[x] = y;
        return true;
    }
};
class Solution {
   public:
    int containVirus(vector<vector<int>>& v) {
        int n = v.size();
        int m = v[0].size();

        auto get_idx = [&](int x, int y) {
            return x * m + y;
        };

        auto get_loc = [&](int idx) -> pii {
            return {idx / m, idx % m};
        };

        unordered_set<int> ok;
        unordered_set<int> not_ok;

        DSU dsu(n * m);
        int d[4][2] = {-1, 0, 1, 0, 0, -1, 0, 1};
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < m; j++) {
                if (v[i][j] == 0) {
                    ok.insert(get_idx(i, j));
                } else {
                    for (int k = 0; k < 4; k++) {
                        int i2 = i + d[k][0], j2 = j + d[k][1];
                        if (i2 >= 0 && i2 < n && j2 >= 0 && j2 < m && v[i2][j2]) {
                            dsu.merge(get_idx(i, j), get_idx(i2, j2));
                        }
                    }
                    not_ok.insert(get_idx(i, j));
                }
            }
        }

        int ans = 0;
        while (not_ok.size() != 0 && ok.size() != 0) {
            // for (int i = 0; i < n; i++) {
            //     for (int j = 0; j < m; j++) {
            //         if (not_ok.count(get_idx(i, j))) {
            //             cout << "1 ";
            //         } else {
            //             cout << "0 ";
            //         }
            //     }
            //     cout << endl;
            // }
            debug(ans);

            unordered_map<int, int> count, use;
            for (auto it : ok) {
                auto [x, y] = get_loc(it);

                unordered_set<int> tmp;
                for (int i = 0; i < 4; i++) {
                    int x2 = x + d[i][0], y2 = y + d[i][1];
                    int idx2 = get_idx(x2, y2);
                    if (x2 >= 0 && x2 < n && y2 >= 0 && y2 < m && not_ok.count(idx2)) {
                        int fidx2 = dsu.leader(idx2);
                        use[fidx2]++;
                        tmp.insert(fidx2);
                    }
                }
                for (auto it : tmp) {
                    count[it]++;
                }
            }

            int fmx = dsu.leader(*not_ok.begin());
            for (auto it : count) {
                if (count[fmx] < it.second) {
                    fmx = it.first;
                }
            }
            ans += use[fmx];

            for (int i = 0; i < n * m; i++) {
                if (dsu.leader(i) == fmx) {
                    not_ok.erase(i);
                }
            }

            vector<int> tmp;
            for (auto it : ok) {
                auto [x, y] = get_loc(it);
                bool flag = true;
                for (int i = 0; i < 4; i++) {
                    int x2 = x + d[i][0], y2 = y + d[i][1];
                    int idx2 = get_idx(x2, y2);
                    if (x2 >= 0 && x2 < n && y2 >= 0 && y2 < m && not_ok.count(idx2)) {
                        dsu.merge(it, idx2);
                        flag = false;
                    }
                }
                if (!flag) {
                    tmp.push_back(it);
                }
            }

            for (auto it : tmp) {
                ok.erase(it);
                not_ok.insert(it);
            }

            for (auto it : not_ok) {
                auto [x, y] = get_loc(it);
                for (int i = 0; i < 4; i++) {
                    int x2 = x + d[i][0], y2 = y + d[i][1];
                    int idx2 = get_idx(x2, y2);
                    if (x2 >= 0 && x2 < n && y2 >= 0 && y2 < m && not_ok.count(idx2)) {
                        dsu.merge(it, idx2);
                    }
                }
            }
        }
        return ans;
    }
};