스택 - 탑문제

Question

1. 현재 학습 진도

3-5 스택부분 수강

2. 어려움을 겪는 부분

백준 탑 문제를 강사님께서 구현해주신 코드로
풀어보고 있는데 강사님 코드를 사용하면
시간초과가 나는 것 같습니다.

3. 시도해보신 내용

N = int(input())
tops = list(map(int, input().split()))


def top_stack(N):
    result = [0] * N
    while tops:
        cur_top = tops.pop()
        for i in range(len(tops) - 1, -1, -1):
            if cur_top <= tops[i]:
                result[len(tops)] = i + 1
                break
                
    print(' '.join(map(str,result)))


top_stack(N)

강사님께서 구현해주신 코드에 입력값을 사용자가 지정하게만 바꿔서 백준문제를 풀어보려고 했는데, 시간초과가 납니다.
제가 혹시 코드에 실수한 부분이 있는지 궁금합니다.

딩코딩코 · Answer

수빈님 오랜만에 질문 탭에서 뵙는 것 같습니다

벌써 겨울이 지나고 봄인데 잘 지내시죠?? 좋은 질문 주셔서 감사합니다!!

강의에서 제시되었던 코드는 반복문에서 스택을 사용하도록 변경한 코드라서, 실제로 큰 입력값을 처리하기에는 효율적이지 않습니다!

현재 코드의 시간복잡도는 O(n²)으로:

각 탑마다 (while heights:)
이전의 모든 탑을 순회하며 확인합니다 (for idx in range(len(heights) - 1, -1, -1):)

백준과 같은 알고리즘 문제 사이트에서는 보다 효율적인 O(n) 알고리즘이 필요합니다. 아래는 스택을 활용한 개선된 코드입니다

def top_stack(N, tops):
    result = [0] * N
    stack = []  # (인덱스, 높이) 저장
    
    for i in range(N):
        # 현재 탑보다 낮은 탑은 스택에서 제거
        while stack and stack[-1][1] < tops[i]:
            stack.pop()
        
        # 스택이 비어있지 않다면, 가장 위의 탑이 수신 탑
        if stack:
            result[i] = stack[-1][0] + 1  # 인덱스는 0부터 시작하므로 +1
        
        # 현재 탑을 스택에 추가
        stack.append((i, tops[i]))
    
    return result

# 백준 제출용 코드
N = int(input())
tops = list(map(int, input().split()))
result = top_stack(N, tops)
print(' '.join(map(str, result)))

이 개선된 알고리즘은

각 탑을 한 번씩만 처리합니다
스택에는 레이저 신호를 수신할 가능성이 있는 탑의 정보만 유지합니다
불필요한 비교를 피해 시간복잡도를 O(n)으로 개선합니다

해당 개선 방식에 대해서 추가적으로 교재에서 다뤄봐도 좋을 것 같습니다!

추후에 알고리즘 영상을 보완하는 시점에 해당 문제 풀이 방법에 대해서도 작성해보도록 하겠습니다

강의 개선 아이디어 주심에 대한 감사의 의미로 커피 기프티콘을 드리겠습니다 아래 카카오톡 오픈 링크로 연락 부탁드립니다!!

https://open.kakao.com/me/ding_coding_co

감사합니다

인프런 커뮤니티 질문&답변