417. 太平洋大西洋水流问题

class Solution:
    def pacificAtlantic(self, heights: List[List[int]]) -> List[List[int]]:
        m = len(heights)
        if m == 0:
            return []

        n = len(heights[0])
        if n == 0:
            return []

        pacific, atlantic = set(), set()

        for i in range(m):
            pacific.add((i, 0))
            atlantic.add((i, n - 1))

        for j in range(n):
            pacific.add((0, j))
            atlantic.add((m - 1, j))

        def extend(pos: set[tuple[int, int]]):
            candidates, visited = {(i, j) for (i, j) in pos}, set()
            while candidates:
                i, j = candidates.pop()
                visited.add((i, j))

                for dx, dy in [(1, 0), (0, 1), (-1, 0), (0, -1)]:
                    x, y = i + dx, j + dy
                    if not (0 <= x < m and 0 <= y < n):
                        continue

                    if heights[x][y] >= heights[i][j]:
                        pos.add((x, y))

                        if (x, y) not in visited:
                            candidates.add((x, y))

        extend(pacific)
        extend(atlantic)

        return [[i, j] for i, j in pacific.intersection(atlantic)]