LC 827 Making a Large Island (DSU)

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

const int dx[4] = {-1, 0, 1, 0};
const int dy[4] = {0, 1, 0, -1};

bool isValid(int r, int c, int n){
    if(r < 0 || r >= n || c < 0 || c >= n){
        return false;
    }
    return true;
}

class DS{
    public:
    map<pair<int, int>, pair<int, int> > parent;
    map<pair<int, int>, int> size;

    DS(){}


    pair<int, int> find(int r, int c){
        pair<int, int> cell = make_pair(r, c);
        if(cell == parent[cell]){
            return cell;
        }
        pair<int, int> pCell = parent[cell];
        return parent[cell] = find(pCell.first, pCell.second);
    }

    void Union(int u1, int v1, int u2, int v2){
        pair<int, int> a = find(u1, v1), b = find(u2, v2);
        if(a != b){
            if(size[a] >= size[b]){
                parent[b] = a;
                size[a] += size[b];
            }
            else{
                parent[a] = b;
                size[b] += size[a];
            }
        }
    }
};


int largestIsland(vector<vector<int> > &grid){
    DS dsu;

    //Init parent and size property of DSU for cell value already 1
    for(int i = 0; i < grid.size(); i++){
        for(int j = 0; j < grid.size(); j++){
            if(grid[i][j] == 1){
                pair<int, int> cell = make_pair(i, j);
                dsu.parent[cell] = cell;
                dsu.size[cell] = 1;
            }
        }
    }

    // Update DSU
    for(int i = 0; i < grid.size(); i++){
        for(int j = 0; j < grid.size(); j++){
            if(grid[i][j] == 1){
                set<pair<int, int> > s;
                for(int k = 0; k < 4; k++){
                    int ni = i + dx[k], nj = j + dy[k];
                    if(isValid(ni, nj, grid.size())){
                        //optonal to only merge cell with unique ultimate parents
                        //As merging with different cells with same ultimate parent will not have any affect on they are merged.
                        // Reducing the size of number of components based on a new connection - there we have to state this condition
                        //Only reduce the number of distinct components when merged with cells that have distinct ultimate parents.
                        if(grid[ni][nj] == 1){
                            pair<int, int> pNewCell = dsu.find(ni, nj);
                            if(s.count(pNewCell) == 0){
                                dsu.Union(i, j, ni, nj);
                                s.insert(pNewCell);
                            }
                        }
                    }
                }
            }
        }
    }

    // //NOT NECCESSARY TO DO A GLOBAL SWEEP
    //1) Ultimate parents always have the correct size immediately after unions.
    //2) DSU invariant: root has the true size, non-roots are irrelevant.
    //Global Sweep to FULLY FLATTEN DSU / PATH COMPRESSED FULLY
    // Doing so, will guarantee that each node in DSU will DIRECTLY point to it's ultimate parent
    for(int i = 0; i < grid.size(); i++){
        for(int j = 0; j < grid.size(); j++){
            if(grid[i][j] == 1){
                pair<int, int> p = dsu.find(i, j);
            }
        }
    }



    int maxSize = 0;
    // For each 0 in grid, see what is the size of new connected component
    for(int i = 0; i < grid.size(); i++){
        for(int j = 0; j < grid.size(); j++){
            if(grid[i][j] == 0){
                set<pair<int, int> > s;
                int currentSize = 0;
                for(int k = 0; k < 4; k++){
                    int ni = i + dx[k], nj = j + dy[k];
                    if(isValid(ni, nj, grid.size())){
                        if(grid[ni][nj] == 1){
                            pair<int, int> upNewCell = dsu.find(ni, nj);
                            if(s.count(upNewCell) == 0){
                                currentSize += dsu.size[upNewCell];
                                s.insert(upNewCell);
                            }
                        }
                    }
                }
                //Add the current cell (if converted to 1) to the whole size
                maxSize = max(maxSize, 1 + currentSize);
            }
            else{
                pair<int, int> pCell = dsu.find(i, j);
                //don't add 1, because the cell is already 1
                maxSize = max(maxSize, dsu.size[pCell]);
            }
        }
    }
    // for(auto p : dsu.parent){
    //     pair<int, int> parP = p.second;
    //     cout << "[" << p.first.first << " " << p.first.second << "] [" << parP.first << " " << parP.second << "]" << '\n';
    // }


    // for(auto p : dsu.size){
    //     cout << "[" << p.first.first << " " << p.first.second << "] = " << p.second << '\n';
    // }
    return maxSize;
}

int main() {
    // your code goes here
    int n;
    cin >> n;
    vector<vector<int> > grid(n, vector<int>(n, 0));
    for(int i = 0; i < n; i++){
        for(int j = 0; j < n; j++){
            cin >> grid[i][j];
        }
    }
    cout << largestIsland(grid) << '\n';
}