[백준] [Python] 6198번 옥상 정원 꾸미기 - 스택

https://www.acmicpc.net/problem/6198

6198번: 옥상 정원 꾸미기

 

 

Python

import sys

arr = [ int(sys.stdin.readline()) for _ in range(int(sys.stdin.readline())) ]

stack = []
ans = [0]*len(arr)

for i in range(len(arr)):
    if stack and stack[-1][0] <= arr[i]:
        for s in stack[:-1]:
            ans[s[1]] = stack[-1][1] - s[1]
        while stack and stack[-1][0] <= arr[i]:
            stack.pop()
    
    stack.append((arr[i],i))   
    #print(stack)

for s in stack[:-1]:
    ans[s[1]] = stack[-1][1] - s[1]      
    
print(sum(ans))

 

 

시간 초과

값을 하나 추가할 때마다

스택을 한 번씩 훑어보며 그 거리를 계산하니 발생한 문제였다.

 

5

5

4

3

2

1   이와 같은 예제에서 n!이나 걸리게 된다

 

그래서

위의 정답 코드를 보면 알다시피

값을 삭제해야하는 첫 순간에만 거리를 계산하도록 바꾸었다.

이 때 마지막 원소가 삽입된 후의 결과가 계산되지 않아

for문 밖에서 한 번 더 계산하는 과정을 추가했다.

import sys

arr = [ int(sys.stdin.readline()) for _ in range(int(sys.stdin.readline())) ]

stack = []
ans = [0]*len(arr)

for i in range(len(arr)):
    while stack and stack[-1][0] <= arr[i]:
        stack.pop()
        
    stack.append((arr[i],i))
    #print(stack)
    for s in stack[:-1]:
        ans[s[1]] = stack[-1][1] - s[1]
    
print(sum(ans))