Cindy's Blog

アマゾンで働いているエンジニアの日常

LeetCode 200. Number of Islands

UnionFind island count direction

https://leetcode.com/problems/number-of-islands/

class Solution {
    public int numIslands(char[][] grid) {
        int rows = grid.length;
        int cols = grid[0].length;
        
        int countIslands = 0;
        int[][] directions = {{0, 1}, {0, -1}, {1, 0}, {-1, 0}};
        UnionFind uf = new UnionFind(rows * cols);
        
        for (int x = 0; x < rows; x++){
            for (int y = 0; y < cols; y++){
                if (grid[x][y] == '1') {
                    countIslands++;
                }
            }
        }
        
        for (int x = 0; x < rows; x++){
            for (int y = 0; y < cols; y++){
                if (grid[x][y] == '1'){
                    for (int[] dir : directions) {
                        int nx = x + dir[0];
                        int ny = y + dir[1];
                        if (nx >= 0 && nx < rows && ny >= 0 && ny < cols && grid[nx][ny] == '1'){
                            if (uf.merge(x * cols + y, nx * cols + ny)){
                                countIslands--;
                            }
                        }
                    }
                }
            }
        }
        
        return countIslands;
        
    }
}

class UnionFind{
    int[] father;
    public UnionFind(int size){
        father = new int[size];
        for (int i = 0; i < size; i++){
            father[i] = i;
        }
    }
    
    public int getFather(int x){
        if (father[x] == x){
            return x;
        }
        return father[x] = getFather(father[x]);
    }
    
    public boolean isSameFather(int a, int b){
        int fa = getFather(a);
        int fb = getFather(b);
        return fa == fb;
    }
    
    public boolean merge(int a, int b){
        int fa = getFather(a);
        int fb = getFather(b);
        if (fa == fb) return false;
        father[fa] = fb;
        return true;
    }
}