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

 

 

※ 해결방법

여기서 오큰수에 대해서 조금 더 자세히 살펴보면,  '오른쪽에 있으면서 Ai보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다.' 라는 개념을 가지고 있는데 왼쪽에서 오른쪽의 방향으로 값을 살펴보면서 기준 인덱스의 크기보다 크면 그 값을 출력해 내면 그 값이 오큰수이다.

 

예시를 통해서 설명하면 3,5,2,7에서 3을 기준으로 놓고 5,2,7 순서로 진행을 하려고 하나 5를 만났을 때 3보다 크므로 바로 5를 출력한다. 그리고 기준을 5로 두고 2,7 순서로 진행을 하는데, 2는 5보다 작으므로 지나치고 7은 5보다 크므로 7을 출력하면 된다. 이런 방식으로 반복하여 값을 출력하고 기준 값보다 큰 값이 없으면 -1을 출력하면된다.

 

하지만, 이 문제는 반복문을 사용하여 문제를 해결하면, 제한 시간을 초과하는 문제를 겪게 된다. (반복문의 시간 복잡도 : O(N^2)

 

시간을 줄이기 위해서 사용할 수 있는 방법은 스택을 사용하는 것이다.

 

스택에서 사용하는 push()와 pop()으로 push()는 다음 인덱스를 확인하는 역할을 하고 pop()은 정답을 정답을 모아두는 배열에 넣는 역할을 한다라고 생각하면 된다.

이 그림처럼 스택에 index를 넣고 list의 값들을 비교한 다음에 만약 오큰수를 만족할 시에, index를 빼내고 그 값에 오큰수 값을 넣어주면 된다.

---

※ 파이썬 코드 정리

n = int(input())
ans = [-1] * n # 오큰수가 없는 부분을 -1로 설정
stack = [] # index를 넣을 리스트
lst = list(map(int, input().split()))

for i in range(n):
    # 스택이 비어있지 않고, 현재 숫자가 스택의 마지막 인덱스에 해당하는 숫자보다 큰 경우
    while stack and lst[stack[-1]] < lst[i]:
        ans[stack.pop()] = lst[i] # 현재 숫자가 다음 큰 숫자가 되므로 ans에 기록
    stack.append(i)
    
while stack:
    ans[stack.pop()] = -1

for i in ans:
    print(i)