Заполнение водой: стек и два указателя

Задача: задержка дождевой воды

height = [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]
# Water trapped at each position:
# pos 2: min(1,3)-0=1
# pos 4: min(2,3)-1=1
# pos 5: min(2,3)-0=2
# pos 6: min(2,3)-1=1
# pos 9: min(3,2)-1=1
# Total = 6
print('height:', height)
print('Expected trapped water: 6')

# Visualise
max_h = max(height)
for row in range(max_h, 0, -1):
    line = ''
    for h in height:
        line += '#' if h >= row else ' '
    print(line)

Подход 1: предварительно вычисленные массивы максимумов

def trap_prefix(height):
    n = len(height)
    if n < 3:
        return 0

    max_left = [0] * n
    max_right = [0] * n

    max_left[0] = height[0]
    for i in range(1, n):
        max_left[i] = max(max_left[i-1], height[i])

    max_right[-1] = height[-1]
    for i in range(n-2, -1, -1):
        max_right[i] = max(max_right[i+1], height[i])

    water = 0
    for i in range(n):
        water += max(0, min(max_left[i], max_right[i]) - height[i])
    return water

print(trap_prefix([0,1,0,2,1,0,1,3,2,1,2,1]))  # 6
print(trap_prefix([4,2,0,3,2,5]))                # 9

Подход 2: два указателя (память O(1))

def trap_two_pointer(height):
    left, right = 0, len(height) - 1
    max_left = max_right = 0
    water = 0

    while left < right:
        if height[left] < height[right]:
            if height[left] >= max_left:
                max_left = height[left]    # new max on the left
            else:
                water += max_left - height[left]  # trapped by max_left
            left += 1
        else:
            if height[right] >= max_right:
                max_right = height[right]
            else:
                water += max_right - height[right]
            right -= 1
    return water

print(trap_two_pointer([0,1,0,2,1,0,1,3,2,1,2,1]))  # 6
print(trap_two_pointer([4,2,0,3,2,5]))                # 9
print(trap_two_pointer([3,0,3]))                      # 3

Почему работают два указателя: инвариант

# Trace two-pointer on [4, 2, 0, 3, 2, 5]
height = [4, 2, 0, 3, 2, 5]
left, right = 0, len(height) - 1
max_l = max_r = water = 0
print('height:', height)
print(f'{'Step':5} {'L':3} {'R':3} {'maxL':5} {'maxR':5} {'water':6} {'total':6}')
step = 0
while left < right:
    side = 'L' if height[left] < height[right] else 'R'
    if side == 'L':
        if height[left] >= max_l: max_l = height[left]
        else:
            w = max_l - height[left]; water += w
        left += 1
    else:
        if height[right] >= max_r: max_r = height[right]
        else:
            w = max_r - height[right]; water += w
        right -= 1
    step += 1
    print(f'{step:5} {left:3} {right:3} {max_l:5} {max_r:5} {water:6}')
print('Total trapped:', water)

Подход 3: монотонный стек (горизонтальные слои)

def trap_stack(height):
    stack = []   # monotonic decreasing indices
    water = 0

    for i in range(len(height)):
        while stack and height[stack[-1]] < height[i]:
            bottom_idx = stack.pop()        # the floor of the valley
            if not stack:
                break                       # no left wall, no water
            left_idx = stack[-1]
            floor = height[bottom_idx]
            water_height = min(height[left_idx], height[i]) - floor
            width = i - left_idx - 1
            water += water_height * width
        stack.append(i)
    return water

print(trap_stack([0,1,0,2,1,0,1,3,2,1,2,1]))  # 6
print(trap_stack([4,2,0,3,2,5]))                # 9

Отслеживание работы монотонного стека

def trap_stack_trace(height):
    stack = []
    water = 0
    for i in range(len(height)):
        print(f'i={i} h={height[i]}: stack={[height[s] for s in stack]}')
        while stack and height[stack[-1]] < height[i]:
            bot = stack.pop()
            if not stack:
                print(f'  Pop {height[bot]}: no left wall, skip')
                break
            left = stack[-1]
            h = min(height[left], height[i]) - height[bot]
            w = i - left - 1
            water += h * w
            print(f'  Pop {height[bot]}: floor={height[bot]}, left_wall={height[left]}, right_wall={height[i]}, h={h}, w={w}, +{h*w}')
        stack.append(i)
    return water

result = trap_stack_trace([0,1,0,2,1,0,1,3,2,1,2,1])
print('Total:', result)

Сравнение всех трёх подходов

height = [0,1,0,2,1,0,1,3,2,1,2,1]

# All three methods — verify they agree
def trap_prefix(h):
    n = len(h)
    ml = [0]*n; mr = [0]*n; ml[0]=h[0]; mr[-1]=h[-1]
    for i in range(1,n): ml[i]=max(ml[i-1],h[i])
    for i in range(n-2,-1,-1): mr[i]=max(mr[i+1],h[i])
    return sum(max(0,min(ml[i],mr[i])-h[i]) for i in range(n))

def trap_two_ptr(h):
    l,r,ml,mr,w = 0,len(h)-1,0,0,0
    while l<r:
        if h[l]<h[r]:
            ml=max(ml,h[l]); w+=ml-h[l]; l+=1
        else:
            mr=max(mr,h[r]); w+=mr-h[r]; r-=1
    return w

def trap_stk(h):
    stk,w = [],[]
    for i in range(len(h)):
        while stk and h[stk[-1]]<h[i]:
            b=stk.pop()
            if not stk: break
            w.append(max(0,min(h[stk[-1]],h[i])-h[b])*(i-stk[-1]-1))
        stk.append(i)
    return sum(w)

for h in [height, [4,2,0,3,2,5], [3,0,3], [1,0,1]]:
    p=trap_prefix(h); t=trap_two_ptr(h); s=trap_stk(h)
    print(f'{h}: prefix={p}, two-ptr={t}, stack={s}, match={p==t==s}')

Контейнер с максимальным количеством воды

def max_water_container(height):
    left, right = 0, len(height) - 1
    max_area = 0

    while left < right:
        area = min(height[left], height[right]) * (right - left)
        max_area = max(max_area, area)
        # Move the shorter bar: moving taller bar can only reduce min
        if height[left] < height[right]:
            left += 1
        else:
            right -= 1
    return max_area

print(max_water_container([1,8,6,2,5,4,8,3,7]))  # 49: bars 8 and 7
print(max_water_container([1,1]))                  # 1
print(max_water_container([4,3,2,1,4]))            # 16

# Key difference from trapping rain water:
# Container: choose 2 bars, water fills freely between them (no internal barriers)
# Trapping:  water fills ALL valleys in the full elevation map

Продвинутый подход: сбор дождевой воды II (3D)

import heapq

def trap_rain_water_2d(heightMap):
    if not heightMap or not heightMap[0]:
        return 0
    m, n = len(heightMap), len(heightMap[0])
    visited = [[False]*n for _ in range(m)]
    heap = []  # (height, row, col)

    # Add all border cells to the heap
    for i in range(m):
        for j in [0, n-1]:
            heapq.heappush(heap, (heightMap[i][j], i, j))
            visited[i][j] = True
    for j in range(n):
        for i in [0, m-1]:
            if not visited[i][j]:
                heapq.heappush(heap, (heightMap[i][j], i, j))
                visited[i][j] = True

    total = 0
    max_h = 0
    while heap:
        h, r, c = heapq.heappop(heap)
        max_h = max(max_h, h)
        for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
            nr, nc = r+dr, c+dc
            if 0<=nr<m and 0<=nc<n and not visited[nr][nc]:
                visited[nr][nc] = True
                total += max(0, max_h - heightMap[nr][nc])
                heapq.heappush(heap, (max(max_h, heightMap[nr][nc]), nr, nc))
    return total

map2d = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]
print(trap_rain_water_2d(map2d))  # 4

Когда использовать каждый метод на собеседованиях

# Quick summary of all three approaches
approaches = [
    {
        'name': 'Prefix max arrays',
        'time': 'O(n)', 'space': 'O(n)',
        'description': '3 passes: build max_left, max_right, sum water column-by-column',
    },
    {
        'name': 'Two pointers',
        'time': 'O(n)', 'space': 'O(1)',
        'description': 'Process smaller side: its max is the limiting wall, no array needed',
    },
    {
        'name': 'Monotonic stack',
        'time': 'O(n)', 'space': 'O(n)',
        'description': 'Compute water in horizontal layers when a taller bar is encountered',
    },
]
for a in approaches:
    print(f'{a["name"]} [{a["time"]} / {a["space"]}]')
    print(f'  {a["description"]}')
    print()

Граничные случаи и распространённые ошибки

def trap(height):
    n = len(height)
    if n < 3:
        return 0   # need at least 3 bars to trap anything

    left, right = 0, n - 1
    max_l = max_r = water = 0
    while left < right:
        if height[left] <= height[right]:
            if height[left] >= max_l:
                max_l = height[left]
            else:
                water += max_l - height[left]  # never negative: max_l > height[left]
            left += 1
        else:
            if height[right] >= max_r:
                max_r = height[right]
            else:
                water += max_r - height[right]
            right -= 1
    return water

# Edge cases
print(trap([]))          # 0: empty
print(trap([1]))         # 0: single bar
print(trap([1,2]))       # 0: two bars
print(trap([3,0,3]))     # 3: simple valley
print(trap([3,3,3]))     # 0: flat top, no water

Быстрая проверка

Итоги урока