[Python] List <빗물 트래핑>

리스트를 사용하는 조금 난이도 있는 문제를 가지고 왔다. 

벽들의 높이를 리스트로 입력받아, 빗물이 쌓일 수 있는 면적을 구하는 문제이다.

처음에는 한개의 포인터로, 더 높은 벽을 만날 때마다 포인터를 높은 벽의 높이로 갱신하며, 이전 포인터 까지의 격차를 더하여 푸려고 계획했다. 

그러나, 포인터를 한개만 사용하면, 가장 높은 벽(예시에서는 3)을 만난 뒤에는 포인터를 갱신할 수 없는 문제가 발생하고, 이 문제를 해결하려면 꽤나 복잡한 로직이 동반된다. 

그래서, 가장 높은 벽을 기준으로 양 끝, 두개의 포인터를 이용하여 풀이를 작성했다.

 

---

1. 두개의 포인터를 이용하여 풀이

from typing import List

def trap(rain: List[int]) -> int:
    sum = 0
    
    left = 0  # left 포인터
    right = len(rain) - 1 # right 포인터
    
    left_max = rain[left] # 포인터 위치 벽의 초기 높이
    right_max = rain[right] # 포인터 위치 벽의 초기 높이

	# left_max와 right_max를 비교하며 더 작은 포인터를 움직이게끔 하면,
    # 가장 높은 벽까지 와서 로직이 멈추게 된다.
    while left < right: 
    	# 현재 포인터 위치의 벽과 초기 높이를 비교하여 갱신 작업을 진행
        left_max = max(left_max, rain[left]) 
        right_max = max(right_max, rain[right])

        if rain[left] <= rain[right]:
         	# 높은 벽의 높이 - 현재 포인터의 높이가 빗물이 쌓일 수 있는 면적이다.
            sum += left_max - rain[left]
            left += 1
        else:
            sum += right_max - rain[right]
            right -= 1

    return sum


if __name__ == '__main__':
    rain = [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]
    print(trap(rain))

위 처럼 두개의 포인터를 이용하여 풀면 O(n)으로 꽤나 효율적인 풀이가 가능하다. 

물론 이렇게 두개의 포인터를 이용하는 방법 말고도 다른 효율적인 방법이 존재한다.

---

2. stack을 이용하여 더 높은 벽(변곡점)을 만날 때마다 stack에 쌓아 놓은 걸 추출하여 계산

from typing import List

def trap(rain: List[int]) -> int:
    stack = []
    sum = 0

    for i in range(len(rain)):
        # stack 에 넣어놓은 마지막 위치의 벽과 현재 벽을 비교
        while stack and rain[i] > rain[stack[-1]]:
            top = stack.pop()

            if not len(stack):
                break

            distance = i - stack[-1] - 1
            waters = min(rain[i], rain[stack[-1]]) - rain[top]

            sum += distance*waters

        stack.append(i)
    return sum


if __name__ == '__main__':
    rain = [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]
    print(trap(rain))

이 방법은 인덱스를 증가시킬 때마다 stack에 인덱스 정보를 넣어두고, 변곡점을 만날 때마다, 빗물의 면적을 계산하여 결과를 도출하는 방식이다.

시간 복잡도는 첫번째 풀이와 비슷하나, 개인적으로는  두개의 포인터를 이용하여 풀었던 방식보다 직관성이 좀 떨어지지 않나 싶다. 필자도 처음에는 이런 식의 풀이를 생각했지만, stack을 꺼내는 방식에서 꽤나 해메서...결국 이 방법을 스스로 생각해낼 수는 없었다.