[baekjoon] 백준 2493번(파이썬): 탑

문제

2493번: 탑

알고리즘

- 반복문을 통해 탑의 높이를 확인한다.

- 스택에 탑이 있는지 확인하고 스택에 탑이 있으면 현재 탑과 비교하고 없으면 스택에 현재 탑의 높이를 추가한다.

- 스택에 탑의 높이를 추가할 때 탑의 위치까지 묶어서 추가하여 탑의 위치를 쉽게 찾게 한다.

- 코드 주석을 확인하면서 보면 더 이해하기 쉬울 것이다.

코드

import sys

n = int(sys.stdin.readline())
top = list(map(int, sys.stdin.readline().split()))
stack = [] # 비교할 탑의 스택
res = [0] * n

# 반복문을 통해 탑의 높이를 확인
for i in range(n):
    # 스택에 탑이 있는지 확인
    if stack:
        while True:
            # 비교해야하는 탑의 높이와 스택에 마지막 탑의 높이를 비교
            # 스택에 마지막 탑의 높이가 더 클 때
            if top[i] <= stack[-1][0]:

                # 현재 탑의 레이저 신호를 수신해야하는 탑이 스택에 마지막 탑이 된다.
                # 스택에 마지막 탑의 위치를 결과 리스트에 + 1 해주어 초기화
                res[i] = stack[-1][1] + 1
                stack.append([top[i], i]) # 현재 비교한 탑의 높이를 추가하고 반복을 멈춘다.
                break

            # 스택에 마지막 탑의 높이가 작다면 마지막 탑을 팝한다.
            # 현재 탑의 높이보다 작은 스택에 마지막 탑은 수신할 수 없기 때문에 팝
            else:
                stack.pop()

            # 스택에 탑이 없을 경우 스택에 현재 높이의 탑을 추가하고 반복을 멈춘다.
            if not stack:
                stack.append([top[i], i])
                break

    # 스택에 탑이 없으면 현재 비교해야하는 탑의 높이와 순서를 스택에 추가
    else:
        stack.append([top[i], i])

print(*res)

 

실패한 코드(시간 초과)

import sys

n = int(sys.stdin.readline())
top = list(map(int, sys.stdin.readline().split()))
stack = [top[0]]
res = [0] * n
for i in range(1, n):

    if stack:
        while True:
            
            if top[i] <= stack[-1]:
                res[i] = top.index(stack[-1]) + 1
                stack.append(top[i])
                break
            else:
                stack.pop()

            if not stack:
                stack.append(top[i])
                break

    else:
        stack.append(top[i])


print(*res)

탑의 수가 1 이상 500.000 이하이기 때문에 시간 복잡도가 O(N^2) 이 넘어가면 안 된다.

그런데 위 코드에서. index() 함수를 사용하면서 시간 복잡도가 O(N^2)이 된 것 같다. 나는 index() 함수의 시간 복잡도가 O(1)인 줄 알았는데 리스트 맨 앞부터 검색해서 인덱스를 찾는 것이기 때문에 O(N)이었다. 이 문제를 해결하기 위해 스택에 탑의 높이를 추가할 때 탑의 위치까지 같이 추가하여 수신하는 탑의 위치를 찾을 수 있게 했다.

github

GitHub - junjange/CodingTest: 내가 푼 코딩 테스트 문제와 해결법