DSA Pacific Atlantic Water Flow

The problem

Pacific Atlantic Water Flow gives a grid of heights. Rain falls on every cell. Water flows from a cell to a side neighbour of equal or lower height.

The Pacific ocean touches the top row and the left column. The Atlantic ocean touches the bottom row and the right column. The answer lists every cell whose water can reach both oceans.

cases = [
    ([[1, 2, 2], [3, 5, 3], [5, 2, 1]], [(0, 1), (0, 2), (1, 1), (1, 2), (2, 0)]),
    ([[2, 1], [1, 2]], [(0, 0), (0, 1), (1, 0), (1, 1)]),
    ([[5, 4, 3], [4, 3, 2], [3, 2, 1]], [(0, 0), (0, 1), (0, 2), (1, 0), (2, 0)]),
]
for heights, expected in cases:
    print("map", heights)
    print("both oceans:", expected)

The third map slopes down toward the Atlantic corner. Water from the inner cells never climbs back to the Pacific edge. Only the five Pacific edge cells qualify there.

Intuition: search downhill from each cell

The first idea starts at one cell and follows the water. A downhill search steps to side neighbours of equal or lower height. Reaching a Pacific edge cell proves one ocean, and reaching an Atlantic edge cell proves the other.

heights = [[1, 2, 2], [3, 5, 3], [5, 2, 1]]
rows, cols = 3, 3
seen = {(1, 1)}
stack = [(1, 1)]
while stack:
    r, c = stack.pop()
    for dr, dc in ((-1, 0), (1, 0), (0, -1), (0, 1)):
        nr, nc = r + dr, c + dc
        if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in seen:
            if heights[nr][nc] <= heights[r][c]:
                seen.add((nr, nc))
                stack.append((nr, nc))
print("cells water from (1, 1) can reach:", sorted(seen))
print("touches Pacific edge:", any(r == 0 or c == 0 for r, c in seen))
print("touches Atlantic edge:", any(r == rows - 1 or c == cols - 1 for r, c in seen))

The centre cell reaches eight of the nine cells, so it touches both edges. The one missed cell is the other height-5 peak, which no downhill step can climb to. The idea works for one cell, and its cost appears when the same search runs again from every other cell.

Brute force: one downhill search per cell

The brute force runs that search once per cell. Each search keeps its own visited set, because equal heights let water loop between cells. A cell joins the answer when its search touches both edges.

def reaches_both(heights):
    rows, cols = len(heights), len(heights[0])
    visits = 0

    def can_reach(sr, sc):
        nonlocal visits
        seen = {(sr, sc)}
        stack = [(sr, sc)]
        pacific = atlantic = False
        while stack:
            r, c = stack.pop()
            visits += 1
            if r == 0 or c == 0:
                pacific = True
            if r == rows - 1 or c == cols - 1:
                atlantic = True
            for dr, dc in ((-1, 0), (1, 0), (0, -1), (0, 1)):
                nr, nc = r + dr, c + dc
                if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in seen:
                    if heights[nr][nc] <= heights[r][c]:
                        seen.add((nr, nc))
                        stack.append((nr, nc))
        return pacific and atlantic

    both = []
    for r in range(rows):
        for c in range(cols):
            if can_reach(r, c):
                both.append((r, c))
    return both, visits

both, visits = reaches_both([[1, 2, 2], [3, 5, 3], [5, 2, 1]])
print("both oceans:", both)
print("cell visits:", visits)

The visit count is the evidence. Nine cells needed far more than nine visits, because neighbouring searches rewalk the same ground. The waste grows with the map.

The optimal solution: climb uphill from both oceans

A reverse search climbs from an ocean to every cell that can flow down to it, and the cells both climbs reach are exactly the answer. That one sentence is the whole method.

Climbing reverses each flow step. Water flows down to equal or lower cells, so the climb steps up to equal or higher cells. Two climbs and one set intersection finish the job.

def both_oceans(heights):
    rows, cols = len(heights), len(heights[0])

    def climb(starts):
        seen = set(starts)
        stack = list(starts)
        while stack:
            r, c = stack.pop()
            for dr, dc in ((-1, 0), (1, 0), (0, -1), (0, 1)):
                nr, nc = r + dr, c + dc
                if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in seen:
                    if heights[nr][nc] >= heights[r][c]:
                        seen.add((nr, nc))
                        stack.append((nr, nc))
        return seen

    pacific = climb([(0, c) for c in range(cols)] + [(r, 0) for r in range(1, rows)])
    atlantic = climb([(rows - 1, c) for c in range(cols)] + [(r, cols - 1) for r in range(rows - 1)])
    return sorted(pacific & atlantic)

cases = [
    ([[1, 2, 2], [3, 5, 3], [5, 2, 1]], [(0, 1), (0, 2), (1, 1), (1, 2), (2, 0)]),
    ([[2, 1], [1, 2]], [(0, 0), (0, 1), (1, 0), (1, 1)]),
    ([[5, 4, 3], [4, 3, 2], [3, 2, 1]], [(0, 0), (0, 1), (0, 2), (1, 0), (2, 0)]),
]
for heights, expected in cases:
    result = both_oceans(heights)
    print("map", heights)
    print("->", result, "match:", result == expected)

All three maps pass. Each climb visits a cell at most once, because the visited set blocks repeats.

Walkthrough: the Pacific climb on the first map

The climb starts on the top row and the left column. The stack holds the frontier. Each step below pops one cell and prints the visited set after its new neighbours join.

heights = [[1, 2, 2], [3, 5, 3], [5, 2, 1]]
rows, cols = 3, 3
starts = [(0, 0), (0, 1), (0, 2), (1, 0), (2, 0)]
seen = set(starts)
stack = list(starts)
print("start:", sorted(seen))
while stack:
    r, c = stack.pop()
    for dr, dc in ((-1, 0), (1, 0), (0, -1), (0, 1)):
        nr, nc = r + dr, c + dc
        if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in seen and heights[nr][nc] >= heights[r][c]:
            seen.add((nr, nc))
            stack.append((nr, nc))
    print("after", (r, c), ":", sorted(seen))

The climb ends with seven cells. Each of them has a downhill path to the Pacific. The Atlantic climb marks seven cells of its own, and the overlap of the two sets is the five cells printed in the optimal section.

Complexity

The brute force runs one search per cell, and each search can visit every cell, so its time is O(cells squared) in the worst case. Each reverse climb visits each cell once, so the optimal time is O(cells).

Both methods store visited sets. The brute force holds one set at a time and the optimal holds two, so both use O(cells) space.

def brute_visits(heights):
    rows, cols = len(heights), len(heights[0])
    visits = 0
    for sr in range(rows):
        for sc in range(cols):
            seen = {(sr, sc)}
            stack = [(sr, sc)]
            while stack:
                r, c = stack.pop()
                visits += 1
                for dr, dc in ((-1, 0), (1, 0), (0, -1), (0, 1)):
                    nr, nc = r + dr, c + dc
                    if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in seen and heights[nr][nc] <= heights[r][c]:
                        seen.add((nr, nc))
                        stack.append((nr, nc))
    return visits

def climb_visits(heights):
    rows, cols = len(heights), len(heights[0])
    total = 0
    for starts in ([(0, c) for c in range(cols)] + [(r, 0) for r in range(1, rows)],
                   [(rows - 1, c) for c in range(cols)] + [(r, cols - 1) for r in range(rows - 1)]):
        seen = set(starts)
        stack = list(starts)
        while stack:
            r, c = stack.pop()
            total += 1
            for dr, dc in ((-1, 0), (1, 0), (0, -1), (0, 1)):
                nr, nc = r + dr, c + dc
                if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in seen and heights[nr][nc] >= heights[r][c]:
                    seen.add((nr, nc))
                    stack.append((nr, nc))
    return total

for n in (2, 3, 4):
    heights = [[(n - r) * n + (n - c) for c in range(n)] for r in range(n)]
    print("size", n, "- brute visits:", brute_visits(heights), " climb visits:", climb_visits(heights))

The printed counts compare the two methods on growing slope maps. The brute force repeats whole searches, and its count pulls away as the map grows. The climbs only ever pay one visit per cell per ocean.

Edge cases and pitfalls

A single cell touches all four edges at once. It reaches both oceans by standing still, so it qualifies.

Equal heights count as flow. A flat map lets water cross every cell, so every cell qualifies.

def both_oceans(heights):
    rows, cols = len(heights), len(heights[0])

    def climb(starts):
        seen = set(starts)
        stack = list(starts)
        while stack:
            r, c = stack.pop()
            for dr, dc in ((-1, 0), (1, 0), (0, -1), (0, 1)):
                nr, nc = r + dr, c + dc
                if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in seen:
                    if heights[nr][nc] >= heights[r][c]:
                        seen.add((nr, nc))
                        stack.append((nr, nc))
        return seen

    pacific = climb([(0, c) for c in range(cols)] + [(r, 0) for r in range(1, rows)])
    atlantic = climb([(rows - 1, c) for c in range(cols)] + [(r, cols - 1) for r in range(rows - 1)])
    return sorted(pacific & atlantic)

print("single cell:", both_oceans([[7]]))
print("flat map:", both_oceans([[2, 2], [2, 2]]))
print("single row:", both_oceans([[3, 1, 2]]))
print("low corner cell:", both_oceans([[2, 2], [2, 1]]))

The last map holds the pitfall. Its low corner cell touches the Atlantic edge, but no downhill path leads it to the Pacific. Touching one ocean is never enough.