<p>You are given an <code>m x n</code> matrix <code>board</code> containing <strong>letters</strong> <code>&#39;X&#39;</code> and <code>&#39;O&#39;</code>, <strong>capture regions</strong> that are <strong>surrounded</strong>:</p>

<ul>
	<li><strong>Connect</strong>: A cell is connected to adjacent cells horizontally or vertically.</li>
	<li><strong>Region</strong>: To form a region <strong>connect every</strong> <code>&#39;O&#39;</code> cell.</li>
	<li><strong>Surround</strong>: The region is surrounded with <code>&#39;X&#39;</code> cells if you can <strong>connect the region </strong>with <code>&#39;X&#39;</code> cells and none of the region cells are on the edge of the <code>board</code>.</li>
</ul>

<p>A <strong>surrounded region is captured</strong> by replacing all <code>&#39;O&#39;</code>s with <code>&#39;X&#39;</code>s in the input matrix <code>board</code>.</p>

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

<div class="example-block">
<p><strong>Input:</strong> <span class="example-io">board = [[&quot;X&quot;,&quot;X&quot;,&quot;X&quot;,&quot;X&quot;],[&quot;X&quot;,&quot;O&quot;,&quot;O&quot;,&quot;X&quot;],[&quot;X&quot;,&quot;X&quot;,&quot;O&quot;,&quot;X&quot;],[&quot;X&quot;,&quot;O&quot;,&quot;X&quot;,&quot;X&quot;]]</span></p>

<p><strong>Output:</strong> <span class="example-io">[[&quot;X&quot;,&quot;X&quot;,&quot;X&quot;,&quot;X&quot;],[&quot;X&quot;,&quot;X&quot;,&quot;X&quot;,&quot;X&quot;],[&quot;X&quot;,&quot;X&quot;,&quot;X&quot;,&quot;X&quot;],[&quot;X&quot;,&quot;O&quot;,&quot;X&quot;,&quot;X&quot;]]</span></p>

<p><strong>Explanation:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2021/02/19/xogrid.jpg" style="width: 367px; height: 158px;" />
<p>In the above diagram, the bottom region is not captured because it is on the edge of the board and cannot be surrounded.</p>
</div>

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

<div class="example-block">
<p><strong>Input:</strong> <span class="example-io">board = [[&quot;X&quot;]]</span></p>

<p><strong>Output:</strong> <span class="example-io">[[&quot;X&quot;]]</span></p>
</div>

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

<ul>
	<li><code>m == board.length</code></li>
	<li><code>n == board[i].length</code></li>
	<li><code>1 &lt;= m, n &lt;= 200</code></li>
	<li><code>board[i][j]</code> is <code>&#39;X&#39;</code> or <code>&#39;O&#39;</code>.</li>
</ul>

from dataclasses import dataclass
from typing import Any

@dataclass
class Node:
    x: int
    y: int
    is_captured: bool
    parent: "Node"
    size: int
    id: int

class UnionFind:
    def __init__(self):
        self.mapping = dict()
    
    def add(self, x, y, is_captured=False):
        if (x, y) not in self.mapping:
            node = Node(x, y, is_captured, None, 1, len(self.mapping))
            node.parent = node
            self.mapping[(x, y)] = node
            return node
        else:
            return self.mapping[(x, y)]

    def find(self, node):
        root = node
        while root.parent != root:
            root = root.parent
        
        while node.parent != root:
            parent = node.parent
            node.parent = root
            node = parent
        
        return root
    
    def union(self, node_a, node_b):
        node_a = self.find(node_a)
        node_b = self.find(node_b)

        if node_a == node_b:
            return
        
        if node_a.size < node_b.size:
            node_a, node_b = node_b, node_a
        
        node_b.parent = node_a
        node_a.size += node_b.size
        node_a.is_captured = node_a.is_captured and node_b.is_captured

class Solution:
    def solve(self, board: List[List[str]]) -> None:
        """
        Do not return anything, modify board in-place instead.
        """
        H = len(board)
        W = len(board[0])

        def neighbors(x: int, y: int):
            if x+1 < W:
                yield x+1, y
            if x-1 >= 0:
                yield x-1, y
            if y+1 < H:
                yield x, y+1
            if y-1 >= 0:
                yield x, y-1
        
        def is_edge(x, y):
            return x == 0 or y == 0 or x == W - 1 or y == H - 1

        sets = UnionFind()

        for y in range(H):
            for x in range(W):
                current = sets.add(x, y, not is_edge(x, y))

                for nx, ny in neighbors(x, y):
                    neighbor = sets.add(nx, ny, not is_edge(nx, ny))

                    if board[y][x] == board[ny][nx]:
                        sets.union(current, neighbor)

        # for y in range(H):
        #     print(" | ".join(f"{sets.find(sets.add(x, y, False)).id: >2}" for x in range(W)))
        
        for y in range(H):
            for x in range(W):
                if sets.find(sets.add(x, y, False)).is_captured:
                    board[y][x] = "X"

28910