[배열] 빗물 트래핑

Trapping Rain Water

투 포인터를 최대로 이용

최대 높이 막대까지 각각 좌우 기둥 최대 높이 left_max, right_max가 현재 높이와의 차이만큼 물 높이 volume을 더해 나간다.
좌우 어느쪽이든 높은 쪽을 향해서 포인터가 가운데로 점점 이동한다.
가장 높이가 높은 막대 지점인 최대 지점에서 좌우 포인터가 서로 만나게 되며 $O(n)$ 풀이가 가능하다.

from typing import List

def trap(height: List[int]) -> int:
    if not height:
        return 0
    volumn = 0
    left, right = 0, len(height) - 1 # 양쪽 끝 투 포인터
    left_max, right_max = height[left], height[right] # 좌우 기둥 높이

	
    while left < right:
    	# 좌우 기둥 최대 높이 갱신
        left_max, right_max = max(height[left], left_max), max(height[right], right_max)

        # 더 높은 쪽을 향해 투 포인터 이동
        if left_max <= right_max:
            volumn += left_max - height[left]
            left += 1
        else:
            volumn += right_max - height[right]
            right -= 1
    return volumn
    

if __name__ == "__main__":
    a = [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2 ,1]
    print(trap(a))

스택 쌓기

스택을 쌓아 나가면서 현재 높이가 이전 높이보다 높을 때, 즉 변곡점(여기서는 기둥이 높아지는 지점인 빨간색 점)를 기준으로 격차만큼 물 높이를 채운다.

이전 높이는 고정된 형태가 아니므로 계속 스택으로 채워나가다가 변곡점을 만날 때마다 하나씩 꺼내면서 이전과의 차이만큼 물 높이를 계산해나간다.

from typing import List


def trap(height: List[int]) -> int:
    stack = []  # 인덱스를 저장하는 스택
    volume = 0

    for i in range(len(height)):
        # 변곡점을 만나는 경우
        while stack and height[i] > height[stack[-1]]:
            top = stack.pop()  # 변곡점 이전 지점 인덱스

            if not len(stack):  # 스택이 비었다면 skip
                break

            # 이전과의 차이만큼 물 높이 처리
            distance = i - stack[-1] - 1  # 이전 기둥과의 거리 계산
            # 변곡 점과 이전 기둥 중 작은 높이 - 변곡점 이전 지점의 높이
            waters = min(height[i], height[stack[-1]]) - height[top]

            volume += distance * waters

            print(
                f"top: {top}, i: {i}, distance: {distance}, waters: {waters}, volume: {volume}, stack: {stack}"
            )
        stack.append(i)
    return volume


if __name__ == "__main__":
    a = [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]
    print(trap(a))

top: 2, i: 3, distance: 1, waters: 1, volume: 1, stack: [1]
top: 5, i: 6, distance: 1, waters: 1, volume: 2, stack: [3, 4]
top: 6, i: 7, distance: 2, waters: 0, volume: 2, stack: [3, 4]
top: 4, i: 7, distance: 3, waters: 1, volume: 5, stack: [3]
top: 9, i: 10, distance: 1, waters: 1, volume: 6, stack: [7, 8]

• 거리와 높이차를 곱해서 파란 선에 해당하는 물의 부피를 구하고 더해 나가는 방법이다.

참고 자료

• 파이썬 알고리즘 인터뷰