[파이썬 알고리즘 인터뷰] 7. 빗물 트래핑

https://github.com/onlybooks/algorithm-interview

 

1. 문제 

leetcode 42. trapping rain water

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

2. 스택을 이용한 내 풀이

 스택에 높이가 낮아지는 곳의 좌표를 저장하는 방법으로 문제를 풀었다. 먼저 주어진 배열의 왼쪽부터 높이를 확인하면서 높이가 낮아진다면 해당 지점의 좌표, 바닥의 높이를 스택에 저장한다. 그러던 중 높이가 높아지는 지점 a가 있다면 스택에서 a보다 높은 위치의 좌표가 나오기 전까지 좌표를 꺼낸다. 그리고 꺼낸 좌표와 a의 좌표, 바닥의 높이를 이용해 추가되는 물의 양을 계산하고 바닥의 높이를 업데이트 한다. 이 방법의 시간복잡도는 O(n)이고, 평균 실행시간은 100ms이다.

def trap(self, height: List[int]) -> int:
    lefts = []
    area = 0
    for i in range(1, len(height)):
        if height[i - 1] > height[i]:
            lefts.append([i, height[i - 1]])
            bottom = height[i]
        elif height[i - 1] < height[i]:
            while len(lefts) > 0:
                if lefts[-1][1] > height[i]:
                    area += (height[i] - bottom) * (i - lefts[-1][0])
                    bottom = height[i]
                    break
                else:
                    left = lefts.pop()
                    area += (left[1] - bottom) * (i - left[0])
                    bottom = left[1]
    return area

 

3. 스택을 이용한 교재의 풀이

 책에서는 내 방법과 다르게 높이가 올라가지 않을 때 모든 높이를 스택에 저장한다. 또한 스택에는 좌표가 아니라 x축 상의 위치만 저장하고, 바닥의 위치도 따로 관리하지 않는다. 왼쪽부터 탐색을 하며 높아지는 지점 a를 만나면 a보다 높은 지점의 좌표가 나오기 전까지 스택에서 위치를 꺼낸다. 그리고 꺼낸 위치와 a의 좌표를 이용해 더할 물의 양을 구한다.시간복잡도는 2번과 같지만, 평균 실행시간은 56ms로 2배 정도 빠르다.

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 len(stack) == 0:
                break
            distance = i - stack[-1] - 1
            water = min(height[i], height[stack[-1]]) - height[top]
            volume += distance * water
        stack.append(i)
    return volume

 

4. 투 포인터를 이용한 교재의 풀이

 이 풀이 방법은 어떤 지점에 물이 들어가는 양은 그 지점의 왼쪽과 오른쪽에서 가장 높은 위치 중 더 낮은 쪽의 높이에 달려있다는 사실을 활용한 방법이다. 이를 이용해서 두개의 포인터가 양쪽 끝부터 left_max와 right_max를 찾아나간다. 두 변수는 가장 높이가 높은 막대를 기준으로 왼쪽과 오른쪽에서 가장 높은 높이를 저장하는 변수이다. 각 포인터가 지나가는 위치 a마다 left_max, right_max 중 작은 값에서 a의 높이를 뺀 값을 물의 양에 더한다. 마찬가지로 O(n)의 시간복잡도를 가지고, 56ms의 시간이 걸린다.

def trap(self, height):
    if not height:
        return 0

    left, right = 0, len(height)-1
    left_max, right_max = height[left], height[right]
    volume = 0
    while left < right:
        left_max, right_max = max(left_max, height[left]), max(right_max, height[right])
        if height[left] <= height[right]:
            volume += left_max - height[left]
            left += 1
        else:
            volume += right_max - height[right]
            right -= 1
    return volume

 지난번 두 수의 합 문제나 이 문제처럼 투포인터가 쓰이는 문제가 생각보담 많은 것 같다. 직관적으로 생각하면 투포인터를 떠올리기 어려운 만큼 문제 푸는 방법을 고민할 때 투포인터, 스택, 큐 등 중에서 어떤 자료구조로 풀 수 있을지 고민해보는게 좋을 것 같다.