Solution to LintCode Number of Islands

NumberOfIslands.cpp

class UnionSet {
    map<int, int> father;
  public:
    UnionSet(set<int> &labels) {
        for(auto label : labels) {
            father[label] = label;
        }
    }
  
    int find(int x) {
        int parent = father[x];
        while(parent != father[parent]) {
            parent = father[parent];
        }
        int temp = -1;
        int fa = x;
        while(fa!=father[fa]) {
            temp = father[fa];
            father[fa] = parent;
            fa = temp;
        }
        return parent;
    }
    
    void set_union(int x, int y) {
        int u_x = find(x);
        int u_y = find(y);
        if(u_x != u_y) {
            father[u_x] = u_y;
        }
    }
};

class Solution {
public:
    /**
     * @param grid a boolean 2D matrix
     * @return an integer
     */
     
    int numIslands(vector<vector<bool>>& grid) {
        
        if(grid.size() == 0 || grid[0].size() == 0) {
            return 0;
        }
        
        int n = grid.size(), m = grid[0].size();
        set<int> lables;
        
        for(int i = 0; i < n; i++) {
            for(int j = 0; j < m; j++) {
                if(grid[i][j] == 1) {
                    lables.insert(i * m + j);
                }
            }
        }
        
        UnionSet us(lables);
        
        for(auto &&label : lables) {
            int dx[] = {0, 0, 1, -1};
            int dy[] = {1, -1, 0, 0};
            for(int i = 0; i < 4; i++) {
                int nextX = label / m + dx[i];
                int nextY = label % m + dy[i];
                if(nextX >= 0 && nextX < n && nextY >=0 && nextY < m) {
                    if(grid[nextX][nextY] == 1) {
                        us.set_union(label, nextX * m + nextY);
                    }
                }
            }
        }
        
        set<int> ans;
        for(auto &&lable : lables) {
            ans.insert(us.find(lable));
        }
        return ans.size();
    }
};