[LeetCode] 42. Trapping Rain Water - 스택 풀이 (Python)

https://leetcode.com/problems/trapping-rain-water/description/

<파이썬 알고리즘 인터뷰> 의 풀이 중 스택 풀이가 잘 이해되지 않아서 정리한 내용을 공유합니다

막대를 하나씩 세운다고 생각했을 때 물을 저장할 수 있는 조건을 생각해 봅시다. 다음 막대로 마지막 막대보다 높이가 같거나 작은 막대가 오면 내리막길 모양이 되어서 물을 채울 수 없습니다. (5, 4, 3, 2, 1)

마지막 막대보다 높이가 높은 막대가 오면 물을 채울 수 있는데, 이 때 이전 막대 중에서 새 막대보다 높이가 낮은 막대 위에만 물을 채울 수 있습니다.

이 때 채울 수 있는 물의 양은 막대 하나와 새 막대의 거리 * min(왼쪽 막대, 오른쪽 막대) 높이 차 로 각 막대와 새 막대의 차이 때문에 담을 수 있는 물의 양을 구합니다.

 

이것을 구하기 위해서 막대의 높이를 Decreasing Monotonic Stack(값이 줄어드는 스택)에 저장합니다. 값이 줄어드는 스택을 만들 수 있으면 값을 스택에 넣고, 만들 수 없으면 값이 줄어드는 스택을 만들 수 있을 때까지 pop 하면서 물의 양을 계산합니다. 막대로 표현하면 내리막길을 만들다가 높은 막대를 만나면 벽이 생기기 때문에 그 막대로 담을 수 있는 물의 양을 계산한다고 생각하시면 됩니다.

실제 구현을 할 때는 distance 때문에 인덱스가 필요하므로 높이 대신 index를 스택에 저장합니다

 

코드

def trap(self, height: List[int]) -> int:
    # Decreasing Monotonic Stack
    # 막대의 index를 저장
    stack = []
    volume = 0

    for i in range(len(height)):
        # 스택의 decreasing 조건을 불만족하면 만족할 때 까지 pop 한다
        # pop 하면서 각 막대로 담을 수 있는 물 높이 계산
        while stack and height[i] > height[stack[-1]]:
            current = stack.pop()

            # 왼쪽 벽 없으면 물 저장 못 함
            if not len(stack):
                break

            left_wall_idx = stack[-1]
            left_wall_height = height[left_wall_idx]

            right_wall_height = height[i]

            distance = i - left_wall_idx - 1
            water_height = min(right_wall_height, left_wall_height) - height[current]

            volume += distance * water_height

        stack.append(i)
    return volume

 

성능

시간 복잡도: O(n)

공간 복잡도: O(n)

 

참고 영상