lakes.py

from typing import List, Tuple

class UnionFind:
    def __init__(self, size: int):
        self.parent = [i for i in range(size)]
        self.rank = [0] * size
        
    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])  # Path compression
        return self.parent[x]
    
    def union(self, x: int, y: int):
        x_root = self.find(x)
        y_root = self.find(y)
        if x_root == y_root:
            return
        # Union by rank
        if self.rank[x_root] < self.rank[y_root]:
            self.parent[x_root] = y_root
        else:
            self.parent[y_root] = x_root
            if self.rank[x_root] == self.rank[y_root]:
                self.rank[x_root] += 1

def count_lakes(grid: List[List[int]]) -> int:
    if not grid:
        return 0
    
    rows, cols = len(grid), len(grid[0])
    uf = UnionFind(rows * cols)
    directions = [(0, 1), (1, 0)]  # Only check right and down to avoid double-counting
    
    # First pass: Build Union-Find structure
    for i in range(rows):
        for j in range(cols):
            if grid[i][j] == 0:  # Only process water cells
                current = i * cols + j
                for di, dj in directions:
                    ni, nj = i + di, j + dj
                    if ni < rows and nj < cols and grid[ni][nj] == 0:
                        neighbor = ni * cols + nj
                        uf.union(current, neighbor)
    
    # Second pass: Count unique lakes
    lakes = set()
    for i in range(rows):
        for j in range(cols):
            if grid[i][j] == 0:
                lakes.add(uf.find(i * cols + j))
    
    return len(lakes)

def lake_size(grid: List[List[int]], coord: Tuple[int, int]) -> int:
    if not grid or not (0 <= coord[0] < len(grid)) or not (0 <= coord[1] < len(grid[0])):
        return 0
    if grid[coord[0]][coord[1]] == 1:  # Land
        return 0
    
    # Reuse the count_lakes Union-Find setup
    rows, cols = len(grid), len(grid[0])
    uf = UnionFind(rows * cols)
    directions = [(0, 1), (1, 0)]
    
    for i in range(rows):
        for j in range(cols):
            if grid[i][j] == 0:
                current = i * cols + j
                for di, dj in directions:
                    ni, nj = i + di, j + dj
                    if ni < rows and nj < cols and grid[ni][nj] == 0:
                        uf.union(current, ni * cols + nj)
    
    target = coord[0] * cols + coord[1]
    root = uf.find(target)
    return sum(1 for i in range(rows) for j in range(cols) 
            if grid[i][j] == 0 and uf.find(i * cols + j) == root)

# Test Cases
grid1 = [
    [1, 1, 0, 1, 1],
    [1, 1, 0, 1, 1],
    [0, 0, 1, 0, 0],
    [1, 1, 0, 1, 1]
]

grid2 = [
    [1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1],
    [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1],
    [1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
    [1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 1, 1],
    [1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1],
    [1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1],
    [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1],
    [1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1]
]

print("Number of lakes in grid1:", count_lakes(grid1))  # Output: 3
print("Size of lake at (2,0) in grid1:", lake_size(grid1, (2, 0)))  # Output: 4
print("Number of lakes in grid2:", count_lakes(grid2))  # Output: 4
print("Size of lake at (3,1) in grid2:", lake_size(grid2, (3, 1)))  # Output: 6


# Time Complexity Analysis

# Initialization:	O(mn)	For grid traversal
# Union-Find:	O(mn α(mn))	α is inverse Ackermann function
# Lake counting:	O(mn)	Final pass through water cells
# Total:	O(mn α(mn))	Effectively linear for practical purposes