[Python] LeetCode 42. Trapping Rain Water

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

문제

풀이

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

class Solution:
    def trap(self, height):
        if not height:
            return 0
        
        volume = 0
        left, right = 0, len(height) - 1
        left_max, right_max = height[left], height[right]
        
        while left < right:
        	if left_max <= right_max:
            	left += 1
                left_max = max(left_max, height[left])
                volume += left_max - height[left]
            right -= 1
            right_max = max(right_max, height[right])
            volume += right_max - height[right]
            
    return volume    	

투 포인터를 이용한 풀이는 left와 right를 양 끝단에 위치시킨 뒤, left_max와 right_max에 각각 height[left], height[right]를 할당하고 시작한다.

left_max(right_max)는 포인터가 움직이면서 경로에 있었던 벽들의 높이 중 제일 높은 값을 저장하는 변수들이다.

코드는 left_max와 right_max의 값 중, 작은 쪽이 큰 쪽으로 포인터를 옮겨가며 해당 포인터 위치의 벽의 높이와 left_max(right_max)를 비교한다. 이 계산 과정은 max(left_max, height[left])(max(right_max, height[right])) 코드에서 이뤄진다.

그 뒤에 left_max(right_max)에서 height[left](height[right])를 빼주면, 해당 포인터 위치의 벽에서 물이 차오를 수 있는 양이 계산이 된다.

이런 식으로 반복문이 계속 돌면 빠지지 않고 모든 공간이 계산된다.

2. 스택을 이용한 풀이

class Solution:
    def trap(self, height):
        stack = []
        volume = 0
        
        for i in range(len(height)):
            while stack and height[i] > height[stack[-1]]:
                top = stack.pop()
                
                if not stack:
                    break
                
                distance = i - stack[-1] - 1
                water = min(height[i], height[stack[-1]]) - height[top]
                volume += distance * water
            stack.append(i)
        return volume

스택을 이용한 풀이는 생각보다 어려웠고 직관적으로 이해하기 힘들었다.

하지만 직접 그림을 그려 흐름을 따라가보려하니 어느정도 이해가 됐기에, 우선 코드를 쭉 살펴본 뒤 그린 그림과 예시를 바탕으로 실제 동작을 살펴보려한다.

...
while stack and height[i] > height[stack[-1]]:
...

위와 같은 while 반복문의 조건은 stack이 비어있지 않고 반복문이 가리키는 현재 벽(height[i])이 stack에 가장 마지막으로 삽입된 벽의 높이보다 높을 때, 즉 벽이 높아지는 변곡점을 만났을 때 를 의미한다.

만약 while 반복문의 조건을 만족하지 않는다면

...
stack.append(i)
...

코드로 건너뛰어 stack에 i의 값이 삽입된다.

while 반복문의 조건을 만족했다면 반복문 내부로 들어가

...
top = stack.pop()
...

을 해주는데, 이 top은 변곡점을 만나기 전에 이미 차오른 물의 양이 계산된 벽을 제외한 벽 중 가장 낮은 벽의 인덱스라고 생각하면 된다.

이 개념이 그림을 그려보기 전에 가장 이해하기 어려웠던 부분이었는데 뒤의 코드들을 마저 살펴보고 밑에서 그림과 함께 설명하겠다.

뒤이어 나오는

...
if not stack:
		break
...

코드는 왼쪽 벽이 존재하지 않아 물이 채워지지 않는 경우에 while 반복문을 마저 돌지 않고 break하게 해주는 부분인데, 나머지 코드들을 살펴본 뒤에 그림으로 설명하겠다.

break된다면 위에서 while 조건문을 만족하지 않았을 때와 마찬가지로

...
stack.append(i)
...

가 실행된다.

...
distance = i - stack[-1] - 1
...

이 코드는 오른쪽 벽(height[i])에서부터 왼쪽 벽(height[stack[-1]])까지의 거리를 나타낸다.

다시 말해 물이 차오르는 공간의 너비라고 생각하면 편하다.

마지막에 1을 한번 빼주는 이유는 왼쪽 벽은 길이에 포함되지 않아야하기 때문이다.

만약 이 1을 빼주지 않는다면 왼쪽 벽까지 길이에 포함하여 계산하게 된다.

...
water = min(height[i], height[stack[-1]]) - height[top]
...

distance가 구해지면 그 뒤에는 물이 차오르는 실질적 높이가 위의 코드를 거치며 계산된다.

water는 오른쪽 벽(height[i])과 왼쪽 벽(height[stack[-1]]) 중 더 낮은 벽의 높이에 맞춰 물이 차오르는 높이를 구한다.

당연히 두 벽 중 낮은 쪽의 높이까지만 물이 차오를 수 있기에 min(height[i], height[stack[-1]])을 통해 낮은 벽의 높이를 구한 뒤, 이미 차오른 물의 양이 계산된 벽 중 가장 낮은 벽의 높이(height[top])를 빼는 것이다.

여기까지의 코드가 동작하면 그 뒤는 volume에 실제 차오른 물의 양을 더해주는

...
volume += distance * water
...

코드가 실행된다.

while 반복문이 한번 돌았으므로, 다시 상기 조건을 충족하는지 확인 후 while 반복문 내부가 실행될지, 아니면 탈출할지가 결정된다.

만약 탈출하게 된다면 마찬가지로

...
stack.append(i)
...

코드를 통해 stack에 삽입되게 된다.

코드의 전체적인 실행 순서 및 흐름을 살펴보았으니 이제 실제 값이 어떤 식으로 할당되는지를 정리해보겠다.

예시로 들 height이다.

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

위의 그림과 같은 경우 변곡점이 총 2개(height[4] → height[5], height[5] → height[6]) 존재할 수 있다.

벽의 높이가 낮아지는 것은 변곡점으로 간주하지 않기 때문에 while 반복문의 조건을 충족시키지 않아 i = 0일 때부터 i = 4일 때 까지는 위에서 코드를 살펴본 것처럼 stack.append(i)만 실행될 것이다.

이제 while 반복문의 조건을 만족하여 핵심 코드가 실행되는 i = 5의 경우를 보자.

해당 시점에서 중요한 변수들의 값은 아래와 같다.

stack = [0, 1, 2, 3, 4]
volume = 0

while 반복문의 조건들

1. not stack: stack = [0, 1, 2, 3, 4]
2. height[i] > height[stack[-1]]: height[5] > height[4] == 4 > 0

모두 만족하므로, while 반복문 내부의 코드가 실행된다.

...
top = stack.pop()
# top = 4
# stack = [0, 1, 2, 3]
...

top에 4이 할당되고, stack = [0, 1, 2, 3]이 된다.

...
if not stack:
# stack = [0, 1, 2, 3]
		break
...

stack은 해당 시점에서 비어있지 않으므로 그대로 진행.

물이 차오르는 실질적인 너비인 distance는 다음과 같이 할당된다.

...
distance = i - stack[-1] - 1
# distance = 5 - 3 - 1
# distance = 1
...

물이 차오르는 높이인 water는 다음과 같이 할당된다.

...
water = min(height[i], height[stack[-1]]) - height[top]
# water = min(height[5], height[3]) - height[4]
# water = min(4, 1) - 0
# water = 1 - 0
# water = 1
...

앞서 설명한 것처럼 왼쪽 벽의 높이(1)와 오른쪽 벽의 높이(4) 중 물이 더 낮은 왼쪽 벽의 높이를 초과해서 차오를 수는 없기 때문에 min(height[i], height[stack[-1]])을 통해 구한 값에서 height[top]을 빼준다.

최종적으로 volume에 더해지는 값은 아래와 같다.

...
volume += distance * water
# volume += 1 * 1
# volume += 1
# volume = 0 + 1 = 1
# volume = 1
...

이 과정을 거쳐 while 반복문 내부의 반복 하나가 끝났다.

계산된 차오른 물의 높이는 다음 그림과 같다.

주요 변수들의 값은 다음과 같다.

stack = [0, 1, 2, 3]
volume = 1

이제 다시 while 반복문의 조건을 충족하는지 살펴보자.

while 반복문의 조건들

1. not stack: stack = [0, 1, 2, 3]
2. height[i] > height[stack[-1]]: height[5] > height[3] == 4 > 1

모두 만족하므로 다시 while 반복문이 실행된다.

...
top = stack.pop()
# top = 3
# stack = [0, 1, 2]
...

top에 3이 할당되고, stack = [0, 1, 2]이 된다.

...
if not stack:
# stack = [0, 1, 2]
		break
...

해당 시점에서 stack은 비어있지 않음.

위에서 이 top은 변곡점을 만나기 전에 이미 차오른 물의 양이 계산된 벽을 제외한 벽 중 가장 낮은 벽의 인덱스라고 설명했다.

이 그림을 잘 살펴보면, 변곡점을 만나기 전 가장 낮은 벽은 height[4]이다.

하지만, 앞선 반복에서 우리는 해당 벽에 차오른 물의 양을 volume에 더해줬기 때문에, 이를 제외한 나머지 벽들 중 가장 낮은 벽은 height[3]에 해당한다.

따라서 top에 3이 할당되는 것이다.

...
distance = i - stack[-1] - 1
# distance = 5 - 2 - 1
# distance = 2
...

distance에는 2가 할당된다.

...
water = min(height[i], height[stack[-1]]) - height[top]
# water = min(height[5], height[2]) - height[3]
# water = min(4, 2) - 1
# water = 2 - 1
# water = 1
...

water에는 1이 할당된다.

volume에 더해지는 값은 다음과 같다.

...
volume += distance * water
# volume += 2 * 1
# volume += 2
# volume = 1 + 2 = 3
# volume = 3
...

또 하나의 반복이 종료됐다.

계산된 차오른 물의 양 부분은 다음과 같다.

주요 변수들의 값은 다음과 같다.

stack = [0, 1, 2]
volume = 3

이제 다시 while 반복문의 조건을 확인할 차례다.

1. not stack: stack = [0, 1, 2]
2. height[i] > height[stack[-1]]: height[5] > height[2] == 4 > 2

아직 모든 조건을 충족시키므로 계속 반복문을 실행시킨다.

...
top = stack.pop()
# top = 2
# stack = [0, 1]
...

...
if not stack:
# stack = [0, 1]
		break
...

...
distance = i - stack[-1] - 1
# distance = 5 - 1 - 1
# distance = 3
...

...
water = min(height[i], height[stack[-1]]) - height[top]
# water = min(height[5], height[1]) - height[2]
# water = min(4, 3) - 2
# water = 3 - 2
# water = 1
...

...
volume += distance * water
# volume += 3 * 1
# volume += 3
# volume = 3 + 3 = 6
# volume = 6
...

반복 종료

계산된 차오른 물의 양 확인

주요 변수들의 값 확인

stack = [0, 1]
volume = 6

다시 while 반복문 조건 확인

1. not stack: stack = [0, 1]
2. height[i] > height[stack[-1]]: height[5] > height[1] == 4 > 3

모든 조건 만족이니, 계속 반복문 실행

...
top = stack.pop()
# top = 1
# stack = [0]
...

...
if not stack:
# stack = [0]
		break
...

...
distance = i - stack[-1] - 1
# distance = 5 - 0 - 1
# distance = 4
...

...
water = min(height[i], height[stack[-1]]) - height[top]
# water = min(height[5], height[0]) - height[1]
# water = min(4, 4) - 3
# water = 4 - 3
# water = 1
...

...
volume += distance * water
# volume += 4 * 1
# volume += 4
# volume = 6 + 4
# volume = 10
...

반복 종료

계산된 차오른 물의 양 확인

주요 변수들의 값 확인

stack = [0]
volume = 10

while 반복문 조건 확인

1. not stack: stack = [0]
2. height[i] > height[stack[-1]]: height[5] > height[0] == 4 > 4

드디어 while 반복문의 조건을 만족시키지 않는 경우가 발생했다.

while 반복문의 조건에 부합하지 않으므로, 코드는 반복문을 실행시키지 않고 바로

...
stack.append(i)
# stack.append(5)
# stack = [0, 5]
...

위의 코드를 실행시켜 stack에 i를 삽입하게 된다.

그 뒤 제일 바깥 for 반복문이 실행되어 i = 6일 때의 반복이 실행되게 된다.

다시 while 반복문의 조건을 만족하는지 살펴보면,

1. not stack: stack = [0, 5]
2. height[i] > height[stack[-1]]: height[6] > height[5] == 5 > 4

모두 만족한다.

따라서 while 반복문이 다시 실행되게 된다.

...
top = stack.pop()
# top = 5
# stack = [0]
...

...
if not stack:
# stack = [0]
		break
...

...
distance = i - stack[-1] - 1
# distance = 6 - 0 - 1
# distance = 5
...

...
water = min(height[i], height[stack[-1]]) - height[top]
# water = min(height[6], height[0]) - height[5]
# water = min(5, 4) - 4
# water = 4 - 4
# water = 0
...

하지만 이 경우, 앞선 반복문에서와는 살짝 다르게 water에 0이 할당된다.

그 이유는, 왼쪽 벽(height[0])의 높이가 변곡점을 만나기 전에 이미 차오른 물의 양이 계산된 벽을 제외한 벽 중 가장 낮은 벽의 높이와 같기 때문에, 이는 물이 차오를 수가 없기 때문이다.

water가 0이기에 volume에 더해지는 값 역시 0이다.

...
volume += distance * water
# volume += 5 * 0
# volume += 0
# volume = 10 + 0
# volume = 10
...

반복 종료

주요 변수들의 값 확인

stack = [0]
volume = 10

while 반복문 조건 확인

1. not stack: stack = [0]
2. height[i] > height[stack[-1]]: height[6] > height[0] == 5 > 4

여전히 조건들에 충족하므로 다시 반복문이 실행된다.

...
top = stack.pop()
# top = 0
# stack = []
...

해당 시점에서의 stack은 여태 봤던 코드와는 다르게 비어있다.

그리고 이 뒤에 실행되는 코드는 다음과 같다.

...
if not stack:
# stack = []
		break
...

앞선 설명에서 이 코드는 왼쪽 벽의 유무를 판단하는 코드라고 설명했다.

stack이 비어있다는 뜻은 아래의 그림과 같이 왼쪽 벽이 존재하지 않는다는 것을 의미한다.

...
stack.append(i)
# stack.append(6)
# stack = [6]
...

그 뒤는 for 반복문의 i 값이 1 증가하여 반복이 시작되는데, 이미 기존 i는 6이었기에 더 이상의 반복은 있을 수 없다.

최종적으로 리턴해야하는 volume의 값을 확인해보면

volume = 10

따라서, 정답은 10이 된다.

느낀 점

처음에는 단순히 정리 목적으로 그림을 그려보며 코드를 분석했는데, 계속 하다보니 스택이라는 자료구조에 대한 이해도도 깊어졌다.
직전의 값들과 현재의 값을 비교해야하는 상황에 스택만큼 이를 편히 다룰 수 있게 해주는 자료구조는 없는 것 같다.