Lintcode 1391. Making A Large Island

In a 2D grid of 0s and 1s, we change at most one 0 to a 1.

After, what is the size of the largest island? (An island is a 4-directionally connected group of 1s).

Example

Example 1:

Input:[[1, 0], [0, 1]]
Output:3

Explanation:
Change one 0 to 1 and connect two 1s, then we get an island with area = 3.

Example 2:

Input: [[1, 1], [1, 0]]
Output:4

Explanation:
 Change the 0 to 1 and make the island bigger, only one island with area = 4.

Example 3:

Input:[[1, 1], [1, 1]]
Output:4

Explanation:
 Can't change any 0 to 1, only one island with area = 4.

Notice

• 1 <= grid.length = grid[0].length <= 50.
• 0 <= grid[i][j] <= 1.

Code (Java):

public class Solution {
    /**
     * @param grid: 
     * @return: nothing
     */
    public int largestIsland(int[][] grid) {
        // 
        if (grid == null || grid.length == 0) {
            return 0;
        }
        
        int m = grid.length;
        int n = grid[0].length;
        
        UF uf = new UF(m, n, grid);
        
        int[][] dirs = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}};
        
        // step 1: connect all 1s
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                for (int[] dir : dirs) {
                    int nx = i + dir[0];
                    int ny = j + dir[1];
                    
                    if (grid[i][j] == 1 && nx >= 0 && nx < m && ny >= 0 && ny < n && grid[nx][ny] == 1) {
                        uf.union(i * n + j, nx * n + ny);
                    }
                }
            }
        }
        
        int maxSize = 0;
        
        // step 2: find all 0s and try to connect
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                int size = 1;
                Set<Integer> parents = new HashSet<>();
                for (int[] dir : dirs) {
                    int nx = i + dir[0];
                    int ny = j + dir[1];
                    
                    if (grid[i][j] == 0 && nx >= 0 && nx < m && ny >=0 && ny < n && grid[nx][ny] == 1) {
                        int parent = uf.getParent(nx * n + ny);
                        if (!parents.contains(parent)) {
                            parents.add(parent);
                            size += uf.getSize(nx * n + ny);
                        }
                    }
                }
                
                maxSize = Math.max(size, maxSize);
            }
        }
        
        return maxSize;
    }
}

class UF {
    int[] parents;
    int[] size;
    int m;
    int n;
    int[][] grid;
    
    public UF (int m, int n, int[][] grid) {
        this.m = m;
        this.n = n;
        this.grid = grid;
        
        parents = new int[m * n];
        size = new int[m * n];
        
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                parents[i * n + j] = i * n + j;
                size[i * n + j] = 1;
            }
        }
    }
    
    public void union(int a, int b) {
        int pa = find(a);
        int pb = find(b);
        
        if (pa != pb) {
            parents[pa] = pb;
            size[pb] += size[pa];
        }
    }
    
    private int find(int a) {
        int root = a;
        
        while (root != parents[root]) {
            root = parents[root];
        }
        
        // path compression
        while (root != a) {
            int parent = parents[a];
            parents[a] = root;
            a = parent;
        }
        
        return root;
    }
    
    public int getSize(int a) {
        int pa = find(a);
        return size[pa];
    }
    
    public int getParent(int a) {
        return find(a);
    }
}