[Baekjoon] 17298번: 오큰수(스택 Gold4) - Python

1. 문제

오큰수

2. 코드

N = int(input())
A = list(map(int, input().split())) // 수열을 저장하는 리스트

answer = [-1] * N
stack = [0]

for i in range(1, N):
    while stack and A[stack[-1]] < A[i]:
        answer[stack.pop()] = A[i]
    stack.append(i)

print(*answer)

3. 로직

1. 오큰수를 저장할 answer 리스트의 모든 원소를 -1로 초기화한다. 오큰수가 없을 경우에 해당 원소의 오큰수는 -1이기 때문이다.
2. 스택을 사용하여 순차적으로 수열을 순회한다. 스택에는 수열의 인덱스를 저장한다. 첫 번째 요소부터 오큰수를 찾기 위해 0번 인덱스로 초기값을 설정한다.
3. 다음 과정을 N-1 번 반복한다.

• 스택이 비어있지 않으면서 A[stack[-1]] < A[i]일 경우:
  • 현재 원소의 오큰수는 해당 원소가 된다.
  • 오큰수를 찾았다는 의미에서 pop() 연산을 통해서 스택에서 해당 인덱스를 제거시킨다.
  • 동시에 answer 리스트의 해당 인덱스에 오큰수인 A[i]를 저장한다.
• 스택이 비어있거나 A[stack[-1]] > A[i]일 경우:
  • 다음 오큰수를 찾기 위해 우선 해당 인덱스를 스택에 추가시킨다.

4. 모든 순회가 끝나면 answer 리스트에 저장된 오큰수들을 출력한다.