[Baekjoon] 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)을 공백으로 구분해 출력한다.

 

 

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

따라서 스택을 이용하여 문제를 풀어 보았다.

 

스택을 활용한 아이디어

1. 스택에 새로 들어오는 수가 top에 존재하는 수보다 크면 그 수는 오큰수이다.

2. 수열에서 오큰수가 존재하지 않는 숫자에 -1을 출력해야 한다.

 

a. 스택이 채워져 있고 A[index] > A[top]인 경우 정답 수열에 오큰수를 저장한다. 여기서 A[index]는 새로 들어온 값이다.

b. 현재 인덱스를 스택에 append 후 다음 인덱스로 넘어간다.

c. a~b를 N만큼 반복한 후 현재 스택에 남아 있는 인덱스에 -1을 저장한다.

 

내 풀이

n = int(input())
A = list(map(int, input().split()))

answer = [0] * n
stack = []

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

for i in range(n):
    result += str(answer[i]) + ''
    
print(result)

 

시간복잡도를 최적화하기 위해 for문과 변수 선언을 최소화 해봤다.

 

n = int(input())
A = list(map(int, input().split()))

ans = [-1] * n
stack = [0]

for i in range(1, n):

    while stack and A[stack[-1]] < A[i]:
        ans[stack.pop()] = A[i]

    stack.append(i)

print(*ans)

 

ans 배열을 0이 아닌 -1로 초기화하여 오큰수가 없는 경우 기본 값을 나타내도록 하였다.

스택도 [0]으로 초기화하여 이후에 비교를 시작할 때 초기 값을 설정해 두었다.

마지막으로 결과를 출력할 때 *을 이용해 리스트의 요소를 개별 인자로 분리해서 출력한다.