[Python] 백준 / 오큰수 / 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)이 주어진다.

4
3 5 2 7

출력

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

5 7 7 -1

접근 방식

1. 수를 num_list에 받고 이중 반복문을 통해 리스트의 현재 값과 그 뒤 리스트를 비교하여 오큰수를 찾는다
   O(n^2) 으로 풀이 가능 => 시간초과

2. stack으로 풀이

   • stack 리스트에 num_list의 인덱스 값을 push하면서 top값보다 큰 값을 만나면 stack.pop()으로 인덱스를 꺼내어 오큰수를 저장한 후 현재 인덱스 값을 다시 push 한다
     -> O(n) 풀이

코드

# O(n^2) 시간초과 풀이
def nge (num_list):
    answer = ''
    for i in range(len(num_list)):
        for j in range(i+1, len(num_list)):
            # 현재 num_list의 값보다 큰 값을 찾으면 오큰수로 설정 후 break
            if num_list[i] < num_list[j]:
                answer+= str(num_list[j]) + ' '
                break
        # 반복문을 정상적으로 통과(오큰수를 찾지 못함)하면 오큰수에 -1 저장
        else :
            answer += '-1 ' 
    
    return answer.rstrip()

_ = input()
num_list = list(map(int, input().split()))

print(range(num_list))

# O(n) 통과 풀이
n = int(input())
num_list = list(map(int, input().split()))

answer = [-1 for _ in range(n)]

stack_ = []

for idx, num in enumerate(num_list):
    if len(stack_) == 0 :
        stack_.append(idx)
    
    elif num_list[stack_[-1]] > num :
        stack_.append(idx)
    
    else :
        while True:
            if len(stack_) != 0 and num_list[stack_[-1]] < num :
                answer[stack_.pop()] = num
            else :
                stack_.append(idx)
                break

print(' '.join([str(i) for i in answer]))