프로그래밍/[ Python ]

[배열] 빗물 트래핑

망나 2021. 1. 13. 06:51

* leetcode 42. Trapping Rain Water

 

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

# 입력
[0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]

# 출력
6

 

높이와 너비 모든 공간을 차례대로 모두 살펴보면 O(n^2)에 풀이가 가능하지만 효율적이지 않기 때문에 투 포인터나 스택을 사용하여 O(n)으로 풀이해야 한다.

 

풀이 1. 투 포인터를 최대로 이동

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_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

 

풀이 2. 스택 쌓기

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