문제
높이를 입력받아 비 온 후 얼마나 많은 물이 쌓일 수 있는지 계산하라.
입력 : height = [0,1,0,2,1,0,1,3,2,1,2,1]
출력 : 6
입력 : height = [4,2,0,3,2,5]
출력 : 9
코드
from typing import List
class Solution:
def trapTwoPointer(self, 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_max, right_max = max(left_max, height[left]), max(right_max, height[right])
if left_max <= right_max:
volume += left_max - height[left]
left += 1
else:
volume += right_max - height[right]
right -= 1
return volume
def trapStack(self, height: List[int]) -> int:
stack = []
volume = 0
for i in range(len(height)):
while stack and height[i] > height[stack[-1]]:
print(i)
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풀이
투포인터
왼쪽 끝에 포인터와 오른쪽 끝에 포인터를 두고, 가장 높은 위치를 향해 서로 좁혀 간다. 왼쪽 max보다 오른쪽 max가 더 크거나 같으면 왼쪽 포인터를 한칸 오른쪽으로 가게 하고, 오른쪽 max보다 왼쪽 max가 더 크다면 오른쪽 포인터를 왼쪽으로 한칸 옮겨 간다. 그러면 가장 높은 위치에서 왼쪽, 오른쪽 포인터가 만나게 된다.
가장 높은 위치로 가는 과정에서 현재의 높이(왼쪽 포인터 or 오른쪽 포인터)와 그 때 까지의 max(왼쪽 or 오른쪽)의 차이 만큼을 계산해서 volume에 더한다.
예를 들면 왼쪽 포인터가 2일 때 왼쪽 max는 1이고 현재의 위치에서 높이는 0이므로 volume에 1이 더해지게 된다. 포인터 2 이후에 벽이 1보다 낮을꺼라는 생각은 안해도 된다. 오른쪽의 max가 왼쪽의 max보다 크거나 같을 때 실행되기 때문이다. 반대로 오른쪽 포인터에서도 마찬가지이다. 아래 표는 포인터가 이동하면서의 max, volume 값의 변화이다.
스택 쌓기
stack에 주어진 높이의 순서인 i를 차례로 넣고, 스택에 쌓인 맨위쪽 위치의 높이와 현재 위치 i의 높이를 비교하여 현재가 더 높을 경우에는 물이 찰수도 있다는 거니깐 stack에서 뽑아서 top변수에 넣는다. 그다음 stack을 확인했을 때 들어가 있는 i가 없다면 왼쪽 부분으로 벽이 없다는거니깐 물이 왼쪽으로 다 샌다고 생각할 수 있으므로 while문을 멈추고 현재의 i를 stack에 넣은뒤 계속 반복한다.
distance = i - stack[-1] - 1
waters = min(height[i], height[stack[-1]]) - height[top]
volume += distance * waters위의 코드는 물이 차 있는 가로(distance), 세로(waters)의 길이를 구해서 곱해서 volume 값에 계속 더해나가는 과정이다. i가 3일 때를 생각해보면 top은 2가되고, 가로 distance는 3-1-1=1이 된다. waters는 min(2,1)-0=1이 되므로 1*1=1을 volume에 더해 volume은 최종 1이 된다. 첫번째 빗물이 담기는 과정이다. distance는 1과 3의 사이의 길이를 구하니깐 3-1-1이라고 할 수 있고, waters는 양쪽 벽의 높이의 최소값과 가운데 물이 들어갈 곳의 높이를 빼준것이라고 할 수 있다. 아래는 for문 및 while문이 진행되면서 stack, top, distance, waters, volume의 변화이다.
출처
파이썬 알고리즘 인터뷰 (글 : 박상길 그림 : 정진호) [책만]