오큰수

오큰수

문제

크기가 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

문제 해석

stack을 이용하는 문제이다. 무조건 리스트의 오른쪽을 비교하라는 문제가 나온다면, stack을 이용하도록 하자. 이 문제의 핵심은 리스트를 탐색하면서 값들을 비교해 보는 것이다. 왜 굳이 stack을 쓰면서 탐색을 할까? 만약 예제 입력 2 처럼 9 5 4 8 이란 숫자가 들어왔다. 그렇다면 9 5 4 까지는 감소하는 부분 수열이기 때문에 오큰수가 나오지 않아서 stack에 계속해서 집어 넣게 된다. 그렇다면 stack에 가장 끝에 쌓여있는 숫자는 가장 작은 숫자가 되는데, 이러한 가장 작은 숫자가 리스트의 다음 수와 비교하여 4라는 숫자가 리스트의 숫자보다 크다면, 나머지 stack에 쌓여있는 숫자들은 비교할 필요도 없게 된다. 반대로 4라는 가장 작은 숫자가 8과 비교하였을 때 8이 더 크면 4를 빼고 다시 순차적으로 5랑 비교를 하고, 9랑 비교를 하는 것이다. 이러한 메커니즘이 stack이다.

풀이

• 리스트의 크기만큼 루프를 돈다.
• stack이 비어있지 않고, stack의 마지막 val값이 리스트의 숫자보다 작다면?
  • 그 리스트의 값이 val의 오큰수이기 때문에 미리 선언해 놓은 ans 리스트에 미리 stack에 넣어놨던 index를 이용하여 오큰수를 저장해준다.
• 비교에 쓰인 리스트의 숫자를 다음 비교를 위해 stack에 index값과 함께 넣어준다.

import sys

N = int(sys.stdin.readline())

NEG = list(map(int, sys.stdin.readline().split()))

stack = []

ans = [-1] * N

for i in range(N):
    while stack and (stack[-1][0] < NEG[i]):
        val, index = stack.pop()
        ans[index] = NEG[i]
    stack.append((NEG[i], i))

print(*ans)

# index 함수는 시간 복잡도가 O(n)이다.

고찰

• 리스트의 index라는 함수는 시간 복잡도가 O(n)이다.
• 처음 풀이는 stack에 굳이 index 값을 튜플 형태로 넣지 않고, 값만 넣은 후, 그 값을 이용해 index 함수에 접근하여 풀이를 하였다.
• 하지만 매번 index 함수를 이용하게 되면 O(n^2) 이여서 완전 탐색과 다를 것이 없다.
• 이러한 연산 속도 때문에 index 함수를 이용하지 않고 stack에 직접 튜플을 이용하여 val와 index를 직접 넣어서 접근하는 것이 연산 속도에 큰 영향을 미친다.