[알고리즘 문제 풀이] 빗물 트래핑

문제

높이를 입력받아 비 온 후 얼마나 많은 물이 쌓일 수 있는지 계산하라.

 

예제1

입력

[0,1,0,2,1,0,1,3,2,1,2,1]

출력

6

---

풀이1 투 포인터를 이용한 풀이

이 문제는 높이와 너비 모든 공간을 차례대로 모두 살펴보면 O(n^2)에 풀이가 가능하다. 하지만 시간 복잡도가 너무 높기 때문에 좀 더 효율적인 풀이를 찾아야 한다. 투 포인터와 스택으로 O(n) 풀이를 할 것이다. 먼저 투 포인터 풀이부터 살펴보자.

이 풀이 방법은 현재 높이와 이전 높이와의 차이만큼 물 높이 volume을 더해 나간다.

이후에는 왼쪽 포인터가 크다면 오른쪽 포인터를 이동시키고, 오른쪽 포인터가 크다면 왼쪽 포인터를 이동시킨다.

투 포인터를 이용한 풀이의 코드는 다음과 같다.

def trap(height: list[int]) -> int:
    if not height:
        return 0

    volume = 0
    left, right = 0, len(height) - 1
    left_max, right_max = height[left], height[right]

    while left < right: # left가 right보다 작은 동안 반복
        left_max, right_max = max(height[left], left_max), max(height[right], right_max)
        # 더 높은 쪽을 향해 투 포인터 이동
        if left_max <= right_max:
            volume += left_max - height[left]
            left += 1
        else:
            volume += right_max - height[right]
            right -= 1
    return volume

height = [0,1,0,2,1,0,1,3,2,1,2,1]
print(trap(height))

---

풀이2 스택을 활용한 풀이

아래 그림과 같이 스택에 쌓아 나가면서 현재 높이가 이전 높이보다 높을 때, 즉 꺾이는 부분 변곡점을 기준으로 격차만큼 물 높이 volume을 채운다. 이전 높이는 고정된 형태가 아니라 들쑥날쑥하기 때문에, 계속 스택으로 채워 나가다가 변곡점을 만날 때마다 스택에서 하나씩 꺼내면서 이전과의 차이만큼 물 높이를 채워 나간다.

스택을 활용한 풀이의 코드는 다음과 같다.

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):
                break

            # 이전과의 차이만큼 물 높이 처리
            distance = i - stack[-1] - 1
            waters = min(height[i], height[stack[-1]]) - height[top]

            volume += distance * waters
        stack.append(i)
    return volume

height = [0,1,0,2,1,0,1,3,2,1,2,1]
print(trap(height))