[Python] 백준 빗물 14719 - Monotonic stack

1. 백준 14719번(빗물)

풀이 방법

모노톤 스택 알고리즘을 활용하여 풀 수 있습니다.

 

Monotonic stack

모노톤 스택 알고리즘은 다음과 같은 방법으로 동작합니다.

1. 스택이 비어 있거나, 스택의 맨 위에 있는 원소가 현재 원소보다 큰 경우, 현재 원소를 스택에 push 합니다.
2. 스택의 맨 위에 있는 원소가 현재 원소보다 작은 경우, 스택에서 원소를 꺼내면서 빗물의 양을 계산합니다. 이 때, 빗물의 양은 스택에서 꺼낸 원소와 현재 원소 사이에 고일 수 있는 빗물의 양입니다.
3. 빗물의 양을 계산한 후, 스택에 현재 원소를 push 합니다.

이 과정을 모든 원소에 대해 반복하면, 스택에는 모노톤한 부분 수열이 저장되고, 빗물의 양을 계산할 수 있습니다.

모노톤 스택 알고리즘은 시간 복잡도가 O(n)으로 매우 효율적이며, 대표적으로 "Trapping Rain Water" 문제와 같은 빗물 문제를 해결하는 데에 사용됩니다.

 

빗물 대표 예제

4 8

3 1 2 3 4 1 1 2 를 시뮬레이션 해보겠습니다.

처음에는 스택에 아무 것도 없으므로, 첫번째 블록들의 갯수인 3을 스택에 넣어줍니다. 더불어 빗물이 담겼을 때 최대 높이가 될 수 있는 max_block 값을 3으로 지정해줍니다.

3다음의 1은 스택의 마지막 원소인 3보다 크지 않으므로 그대로 스택에 넣어줍니다.

세번째 블록갯수인 2와 스택의 마지막 원소인 1을 비교해보았을 때 2가 더 크고, max_block인 3보다는 작기 때문에

1블록 위에는 2-1stack.pop() 만큼의 빗물이 고일 수 있습니다. (정확하게 고인 빗물을 측정하기 위해서 max_block과 현재 block 비교해서 더 작은 값에서 스택 마지막 요소를 빼주어야 합니다.)

answer에 빗물(2-stack.pop())을 더해줍니다. 이때 pop한 횟수 또한 cnt 변수 등을 통해 기록해줍니다.

여기서부터 중요한 점은, 스택의 마지막 원소를 cnt만큼 pop 해주었으니 다시 cnt만큼 세번째 블록의 갯수(2)를 스택에 누적합니다.

위의 사례에서는 pop을 한번 하였으므로, 한번만 누적합니다.

위의 상황에서 3번째 블록 개수인 2을 또 스택에 넣어주면 첫번째, 두번째, 세번째 블록의 상태와 스택의 상태가 순서대로 똑같이 됩니다.(3->2->1)

그 다음 3이 올 때, 3은(max_block과 현재 block 비교해서 더 작은 값에서 스택 마지막 요소를 빼주어야 하는데 똑같으니, 그냥 3) 최대높이를 넘지 않으면서 앞의 두 라인보다 높기 때문에
3-(stack.pop())만큼의 빗물을 answer에 빗물을 더해주고,

 

역시 3-(stack.pop())만큼의 빗물을 answer에 더해줍니다. 

 

이제 다시 cnt 만큼 3을 스택에 다시 넣어주고, 4번째 원소인 3도 넣어줍니다.

첫번째, 두번째, 세번째, 네번째 블록의 상태와 스택의 상태가 순서대로 똑같이 됩니다.(3->3->3->3)

 

다음은 블록 4개입니다! 4는 max_block인 3보다 크기 때문에 

4보다 작은 스택들은 없습니다. 따라서 스택이 비워질 때 까지 answer += 3 - stack.pop()을 해주게 되는데, 사실상 0을 계속 더해주는 효과가 납니다. 그렇게 스택을 비워낸 횟수만큼 아래 그림과 같이 4를 스택에 넣어줍니다.

아래부터는 위에서 서술한 방식과 동일하므로 시뮬레이션 그림만 남겨두도록 하겠습니다.

 

 

전체 코드

import sys
input = sys.stdin.readline
H,W = map(int,input().split())
blocks = [int(i) for i in input().split()]
stack = []
answer = 0
max_val = blocks[0]
for i in range(W):
    cnt = 0
    while stack and stack[-1] < blocks[i]:
        answer += min(blocks[i],max_val) - stack.pop()
        cnt += 1
    if cnt:
        for j in range(cnt):
            stack.append(blocks[i])
    
    if max_val < blocks[i]:
        max_val = blocks[i]
    stack.append(blocks[i])
print(answer)

만약 스택의 크기를 작게 관리하고 싶다면, 아래와 같은 코드로 작성하여도 정답처리 됩니다.

아래 코드 스타일로 관리하였을 때 최종 스택의 모습

import sys
input = sys.stdin.readline
H,W = map(int,input().split())
blocks = [int(i) for i in input().split()]

stack = []
answer = 0
max_height = 0
for i in range(W):
    cnt = 0
    while stack and stack[-1] < blocks[i]: 
        if blocks[i] > max_height:
            answer += max_height - stack.pop()
        else: 
            answer += blocks[i]-stack.pop()
            cnt += 1
    if cnt:
        for j in range(cnt):
            stack.append(blocks[i])

    max_height = max(max_height, blocks[i])
    stack.append(blocks[i])
print(answer)