알고리즘

[leetcode/python] 42번 Trapping Rain Water

grin-quokka 2023. 4. 13. 10:06

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

요약 : 벽 사이에 고인 물의 양을 계산하기

Hard 단계의 문제, 너무 어려웠다.

기본 로직은 다음과 같다

  1. height 리스트를 돈다
  2. 스택에 0번째 값이 현재 벽 높이 보다 크거나 같으면 물이 고인다
    1. 스택을 pop하면서 물 높이를 더해준다
  3. 스택에 벽 높이를 쌓는다
  4. 마지막까지 스택이 비어있지 않다면
  5. 뒤에서부터 벽 높이를 정해서 고인 물을 더해준다 

 

class Solution:
    def trap(self, height: List[int]) -> int:
        result = 0
        stk = []
        for h in height:
            if stk and stk[0] <= h:
                # 벽 높이 기준은 stk[0]으로 나머지 물 높이 구하기
                i = len(stk) - 1
                max = stk[0]
                while i >= 0:
                    water = max - stk.pop()
                    if water > 0:
                        result += water
                    i -= 1
            stk.append(h)
        # 마지막에 stk에 남아있다면
        max = 0
        for i in range(len(stk) - 1, 0, -1):
            if stk[i] > max:
                max = stk[i]
            else:
                result += max - stk[i]

        return result

 

시간복잡도 : O(n)