[🥇4 / 백준 17298 / 파이썬] 오큰수

 

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이다.

입력

첫째 줄에 수열 A의 크기 N (1 ≤ N ≤ 1,000,000)이 주어진다. 둘째에 수열 A의 원소 A1, A2, ..., AN (1 ≤ Ai ≤ 1,000,000)이 주어진다.

출력

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

 

예제 입력 1  4 3 5 2 7	예제 출력 1  5 7 7 -1
예제 입력 2  4 9 5 4 8	예제 출력 2  -1 8 8 -1

 

 

풀이

 

스택을 어떤 방식으로 이용해야할 지 많이 고민한 문제이다.

동적 계획법 문제랑 비슷해 보이기도 해서 더 고민됐던 것 같다.

 

 

모든 수의 오큰수를 찾아야 하기 때문에 아직 오큰수를 찾지 못한 수를 담는 스택을 하나 만들자.

그리고 수열을 순환하며 스택이 비어있을 경우, 자신의 인덱스를 스택에 담고 순서를 넘기도록 한다.

 

현재 비교하는 수열의 수가 스택의 top에 해당하는 인덱스의 수보다 크다면,

스택의 원소를 모두  pop 하면서 현재 수로 바꿔준다.

 

이해하기 쉽게 예제를 보자.

 

왼쪽은 수열(seq), 오른쪽은 스택(stack)이다.

 

[3, 5, 2, 7] [ ]

 

[3, 5, 2, 7] [0]     

스택이 비었으므로  0 인덱스를 가져오고 순서를 넘긴다.

 

[5, 5, 2, 7] [1]     

현재 수 5는 seq[stack[-1]] 즉, seq[0] 보다 크므로 스택의 원소를 pop 하고 5로 바꾸어준다.

그리고 스택이 빌 것이므로 현재 인덱스를 push 해준다.

 

[5, 5, 2, 7] [1, 2]

현재 수 2는 seq[stack[-1]] 즉, seq[1] 보다 작으므로 현재 인덱스를 스택에 push 해준다.

 

[5, 5, 7, 7] [1]

현재 수 7은 seq[2] 보다 크므로  스택의 원소를 pop 하고 7로 바꾸어준다.

 

[5, 7, 7, 7] [3]

인덱스 1 까지 pop 해주고 7로 바꾼 뒤, 스택이 비었으므로 현재 인덱스인 3을 추가해준다.

 

이제 수열의 모든 원소를 순환한 뒤 스택에 3이 남았다.

수열을 모두 순환했는데도 스택에 인덱스가 남아있다는 뜻은 해당하는 수들의 오큰수가 없다는 의미이다.

 

즉, 수열 순환을 마친 뒤 남은 스택의 인덱스에 해당하는 수들을 모두 -1로 바꾸어주면 된다.

 

따라서 최종적으로 [5, 7, 7, -1] 이 된다.

 

아래는 위의 알고리즘을 눈으로 볼 수 있는 코드이다.

 

n = int(input())
seq = list(map(int,input().split()))
stack = []

for i in range(n):
    while(stack):
        print(seq,stack)
        if seq[i] > seq[stack[-1]]:
            seq[stack.pop()] = seq[i]
        else:
            stack.append(i)
            break
    
    if not stack:
        stack.append(i)
        
print(seq,stack)
for i in stack:
    seq[i] = -1
    
print(*seq)

 

 

 

 

정답 코드

 

import sys
input = sys.stdin.readline

n = int(input())
seq = list(map(int,input().split()))
stack = []

for i in range(n):
    while(stack):
        if seq[i] > seq[stack[-1]]:
            seq[stack.pop()] = seq[i]
        else:
            stack.append(i)
            break
    
    if not stack:
        stack.append(i)
for i in stack:
    seq[i] = -1
    
print(*seq)