<p>You are given an <code>n x n</code> binary matrix <code>grid</code>. You are allowed to change <strong>at most one</strong> <code>0</code> to be <code>1</code>.</p>

<p>Return <em>the size of the largest <strong>island</strong> in</em> <code>grid</code> <em>after applying this operation</em>.</p>

<p>An <strong>island</strong> is a 4-directionally connected group of <code>1</code>s.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> grid = [[1,0],[0,1]]
<strong>Output:</strong> 3
<strong>Explanation:</strong> Change one 0 to 1 and connect two 1s, then we get an island with area = 3.
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> grid = [[1,1],[1,0]]
<strong>Output:</strong> 4
<strong>Explanation: </strong>Change the 0 to 1 and make the island bigger, only one island with area = 4.</pre>

<p><strong class="example">Example 3:</strong></p>

<pre>
<strong>Input:</strong> grid = [[1,1],[1,1]]
<strong>Output:</strong> 4
<strong>Explanation:</strong> Can&#39;t change any 0 to 1, only one island with area = 4.
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>n == grid.length</code></li>
	<li><code>n == grid[i].length</code></li>
	<li><code>1 &lt;= n &lt;= 500</code></li>
	<li><code>grid[i][j]</code> is either <code>0</code> or <code>1</code>.</li>
</ul>

from collections import defaultdict, deque

class UnionFind:
    def __init__(self):
        self.parent = {}
        self.sizes = {}

    def add(self, i):
        self.parent[i] = i
        self.sizes[i] = 1

    def find(self, i):
        # If i is parent of itself
        if self.parent[i] == i:
            # Then i is the parent
            return i
        result = self.find(self.parent[i])
        # Cache result
        self.parent[i] = result

        return result
    
    def union(self, i, j):
        i = self.find(i)
        j = self.find(j)
        if i == j:
            return
        
        if self.sizes[i] < self.sizes[j]:
            i, j = j, i
        self.sizes[i] += self.sizes[j]
        self.parent[j] = i

class Solution:
    def largestIsland(self, grid: List[List[int]]) -> int:
        max_rc = len(grid) - 1

        uf = UnionFind()
        def uf_idx(r, c):
            return r * len(grid) + c

        for r in range(len(grid)):
            for c in range(len(grid)):
                island_id = grid[r][c]
                if island_id == 0:
                    continue

                this_uf_idx = uf_idx(r,c)
                uf.add(this_uf_idx)
                if r > 0 and grid[r-1][c] == 1:
                    above_uf_idx = uf_idx(r-1,c)
                    uf.union(this_uf_idx, above_uf_idx)
                if c > 0 and grid[r][c-1] == 1:
                    left_uf_idx = uf_idx(r,c-1)
                    uf.union(this_uf_idx, left_uf_idx)

        max_island = 0
        if len(uf.sizes) > 0:
            max_island = max(uf.sizes.values())
            
        for r in range(len(grid)):
            for c in range(len(grid)):
                val = grid[r][c]
                if val != 0:
                    continue
                
                this_uf_idx = uf_idx(r, c)
                neighbour_islands = set()
                neighbours = [
                    (r - 1, c), 
                    (r, c + 1), 
                    (r + 1, c), 
                    (r, c - 1)
                ]
                for r, c in neighbours:
                    if r < 0 or r > max_rc or c < 0 or c > max_rc:
                        continue
                    neighbour_island_id = grid[r][c]
                    if neighbour_island_id != 0:
                        neighbour_uf_idx = uf_idx(r, c)
                        parent = uf.find(neighbour_uf_idx)
                        neighbour_islands.add(parent)
                combined_island_size = sum([uf.sizes[n] for n in neighbour_islands]) + 1
                max_island = max(max_island, combined_island_size)

        return max_island

# 1 0 1 0 1
# 0 1 1 0 1
# 1 1 1 0 0
# 1 0 1 1 1
# 0 0 1 1 0

# 2 0 3 0 4
# 0 5 3 0 4
# 6 5 3 0 0
# 6 0 3 3 3
# 0 0 3 3 0

40356