<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 enum import Enum

class Position(NamedTuple):
    x: int
    y: int

class UnionFind():
    def __init__(self, positions: List[Position]):
        self.parent = {}
        self.data = {}
        self.rank = {}
        self.size = {}
        self.id = {}

        for idx, position in enumerate(positions):
            self.parent[idx] = idx
            self.data[idx] = position
            self.rank[idx] = 1
            self.size[idx] = 1
            self.id[position] = idx
    
    def find(self, x: int) -> int:
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]

        return x

    def union(self, x: int, y: int):
        x_root = self.find(x)
        y_root = self.find(y)

        if x_root == y_root:
            return

        x_rank = self.rank[x_root]
        y_rank = self.rank[y_root]

        if x_rank > y_rank:
            self.size[x_root] += self.size[y_root]
            self.parent[y_root] = x_root
        
        elif y_rank > x_rank:
            self.size[y_root] += self.size[x_root]
            self.parent[x_root] = y_root
        
        elif x_rank == y_rank:
            self.size[x_root] += self.size[y_root]
            self.parent[y_root] = x_root
            self.rank[x_root] += 1


class TerrainType(Enum):
    WATER = 0
    LAND = 1

class Solution:
    def get_positions(self, grid: List[List[int]], terrain_type: TerrainType) -> List[Position]:
        num_rows = len(grid)
        num_cols = len(grid[0])

        selected_terrain = []
        for i in range(num_rows):
            for j in range(num_cols):
                if TerrainType(grid[i][j]) == terrain_type:
                    new_pos = Position(i, j)
                    selected_terrain.append(new_pos)
        
        return selected_terrain
    
    def is_position_valid(self, position: Position, grid: List[List[int]]) -> bool:
        num_rows = len(grid)
        num_cols = len(grid[0])
        x, y = position
        return x >= 0 and x < num_rows and y >= 0 and y < num_cols

    def create_islands(self, grid: List[List[int]], uf: UnionFind):
        num_rows = len(grid)
        num_cols = len(grid[0])

        directions = [
            Position(0, 1),
            Position(1, 0),
            Position(0, -1),
            Position(-1, 0),
        ]

        for i in range(num_rows):
            for j in range(num_cols):
                curr_pos = Position(i, j)
                if TerrainType(grid[i][j]) == TerrainType.LAND:
                    for direction in directions:
                        di, dj = direction
                        new_i = i + di
                        new_j = j + dj
                        new_pos = Position(new_i, new_j)

                        if self.is_position_valid(new_pos, grid) and TerrainType(grid[new_i][new_j]) == TerrainType.LAND:
                            x_idx = uf.id[curr_pos]
                            y_idx = uf.id[new_pos]

                            uf.union(x_idx, y_idx)
    
    def get_islands(self, uf: UnionFind) -> List[int]:
        island_roots = []
        for idx in uf.parent:
            if uf.parent[idx] == idx:
                island_roots.append(idx)

        return island_roots

    def water_to_land(self, grid: List[List[int]], water: List[Position], uf: UnionFind ,max_island_size) -> int:
        directions = [
            Position(0, 1),
            Position(1, 0),
            Position(0, -1),
            Position(-1, 0),
        ]

        for pos in water:
            adj_positions = []
            for direction in directions:
                x, y = pos
                dx, dy = direction
                
                new_x = x + dx
                new_y = y + dy

                new_pos = Position(new_x, new_y)

                if self.is_position_valid(new_pos, grid) and TerrainType(grid[new_x][new_y]) == TerrainType.LAND:
                    adj_positions.append(new_pos)
            
            unique_roots = set()
            for pos in adj_positions:
                pos_idx = uf.id[pos]
                pos_root = uf.find(pos_idx)
                unique_roots.add(pos_root)
            
            total_size = 1
            for root in unique_roots:
                total_size += uf.size[root]
            
            max_island_size = max(max_island_size, total_size)
        
        return max_island_size
    
    def largestIsland(self, grid: List[List[int]]) -> int:
        land = self.get_positions(grid, TerrainType.LAND)
        uf = UnionFind(land)
        self.create_islands(grid, uf)
        islands = self.get_islands(uf)
        
        island_sizes = [uf.size[island] for island in islands]

        if island_sizes:
            max_island_size = max(island_sizes)
        else:
            max_island_size = 0
        
        water = self.get_positions(grid, TerrainType.WATER)

        return self.water_to_land(grid, water, uf, max_island_size)

77406