[문제 풀이] 오큰수 구하기

2023.08.04 - [알고리즘] - [Do it! 강의 정리] 자료구조에서 스택과 큐와 관련된 문제 

[Do it! 강의 정리] 자료구조

 

백준 온라인 저지 17298번

크기가 N인 수열 A = A1, A2, ..., AN이 있다. 수열의 각 원소 Ai에 관련된 오큰수 NGE(i)를 구하려고 한다. Ai의 오큰수는 오른쪽에 있으면서 Ai보다 큰 수 중 가장 왼쪽에 있는 수를 의미한다. 이러한 수가 없을 때 오큰수는 -1이다. 예를 들어 A = [3,5,2,7]일 때 NGE(1) = 5, NGE(2) = 7, NGE(3) = 7, NGE(4) = -1이다.  A = [9,5,4,8]일 때 NGE(1) = -1, NGE(2) = 8, NGE(3) = 8, NGE(4) = -1이다.

입력
1번째 줄에 수열 A의 크기 N(1<=N<1,000,000), 2번째 줄에 수열 A의 원소 A1, A2, ..., AN(1<=Ai<=1,000,000)이 주어진다.

출력
총 N개의 수 NGE(1), ... , NGE(N)을 공백으로 구분해 출력한다.

문제 분석하기

N의 최대 크기가 1,000,000이므로 반복문으로 오큰수를 찾으면 제한시간을 초과한다.

따라서 스택을 사용하여 다음 아이디어를 추가해 문제를 풀어본다.

 

핵심 아이디어

• 스택에 새로 들어오는 수가 top에 존재하는 수보다 크면 그 수는 오큰수가 된다.
• 오큰수를 구한 후 수열에서 오큰수가 존재하지 않는 숫자에 -1을 출력해야 한다.

문제 푸는 순서

1. 스택이 채워져있거나 A[idex] > A[top]인 경우 pop한 인덱스를 이용하여 정답 수열에 오큰수를 저장함
   pop은 조건을 만족하는 동안 계속 반복
2. 현재 인덱스를 스택에 push(append)하고 다음 인덱스로 넘어감
3. 과정 1, 2를 수열 길이만큼 반복한 다음 현재 스택에 남아있는 인덱스에 -1을 저장

정답

n = int(input())

a = list(map(int, input().split()))
stack = []
result = [0] * n

for i in range(n):
    while stack and a[stack[-1]] < a[i]: #오큰수 조건
        result[stack.pop()] = a[i] # 정답 리스트에 오큰수 저장
    stack.append(i)
while stack:
    result[stack.pop()] =-1
for i in result:
    print(i, end=" ")