[Python] 백준 2493 - 탑 (스택)

제약 사항

• 시간 제한 : 1.5 초
• 메모리 제한 : 128MB

문제

입력

출력

예제 입력

5
6 9 5 7 4

예제 출력

0 0 2 2 4

잘못된 접근 방식

import sys
input = sys.stdin.readline

n = int(input())
li = list(map(int, input().rstrip().split()))
result = []

for i in range(len(li)-1, -1, -1):
    temp = li[i]
    st = []

    if i == 0: # 맨 앞의 탑인 경우
        result.append(0)
    else:
        for j in range(i-1, -1, -1): # 이전 탑들의 높이 탐색
            if li[j] > temp and not st:
                st.append((li[j], j+1)) # (탑 높이, 탑 위치)
                break

        if st:
            height, index = st.pop()
            result.append(index)
        else:
            result.append(0)

for r in result[::-1]:
    print(r, end=' ')

맨 뒤에서부터 완전 탐색 방식으로 구현하였으나 시간초과가 발생하였다. 탑의 개수가 최대 50만개 이므로 브루트 포스 알고리즘을 이용하면 풀 수 없다는 것을 알았다.

올바른 접근 방식

이중 반복문 방식이 아닌 한 번의 반복을 통해서 답을 찾을 수 있는 방법을 고민했다. 내가 도출한 접근 방식은 스택을 이용하여 왼쪽 데이터부터 값을 스택에 추가하며 필요 없는 탑의 요소는 제거하는 방식이다.

 

- st 변수는 (인덱스, 탑의 높이) 형식의 데이터로 구성, 예) [(0,6), (1,9)]

- result 리스트는 수신하는 탑의 인덱스를 저장하는 변수, 예) [0, 0, 2, 2, 4]

 

• 스택이 비어있으면 수신할 수 있는 탑이 없다는 의미이므로 결과를 출력할 리스트에는(result) 0을 추가
• 스택이 비어있지 않다면 스택에 들어있는 탑의 높이들과 현재 탑이 높이를 비교
  • 현재 탑의 높이 < 스택 탑의 높이인 경우, 스택에 있는 탑이 수신할 수 있다는 의미이므로 비교한 탑의 인덱스를 가져온다.
    그리고 result 리스트에 가져온 인덱스+1 를 추가한 뒤 반복문을 종료한다.
  • 현재 탑의 높이 > 스택 탑의 높이인 경우, 스택에 있는 탑이 수신할 수 없다는 의미이므로 스택에서 해당 탑을 pop 한다.
    예를 들어 스택에 [(0, 6)] 이 있고 현재 탑의 높이가 9라면, 높이가 6인 탑은 수신할 수 없으므로 (0,6)을 스택에서 제거한다.
  • 스택이 비어있는 경우가 되거나 조건에 만족하는 탑을 찾을 때까지 높이 비교 로직을 반복한다.
• 스택 비교 반복문이 끝나면 현재 탑의 위치와 높이를 스택에 추가한다.

import sys
input = sys.stdin.readline
n = int(input())
li = list(map(int, input().rstrip().split()))

st = []
result = []

for i in range(n):
    while st:
        current_top_height = li[i] # 현재 탑 높이

        # 현재 탑의 높이보다 스택에 있는 탑의 높이가 낮은 경우
        if st[-1][1] < current_top_height:
            st.pop()
        # 현재 탑의 높이보다 스택에 있는 탑의 높이가 높은 경우
        else:
            result.append(st[-1][0] + 1)
            break
            
    if not st: # 스택이 비어있는 경우
        result.append(0)
    
    st.append((i, li[i])) # (인덱스, 탑의 높이) 추가
    
print(*result)

---

문제 출처

• https://www.acmicpc.net/problem/2493