오큰수

조소복·2022년 5월 27일

문제

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

문제 풀이

시간 초과한 풀이

import sys
input=sys.stdin.readline

N=int(input())
sequence=list(map(int,input().split()))

for i in range(len(sequence)):
    num=sequence.pop(0)
    if len(sequence)!=0 and max(sequence)>num:
        for s in sequence:
            if num<s:
                print(s,end=' ')
                break
    else:
        print(-1,end=' ')

시간초과가 발생할 것 같았지만 확인하기 위해 for문을 중첩하여 하나씩 값을 확인하고 출력하는 방식을 사용했다.

당연히 시간초과가 발생했고 문제를 푸는 방법을 찾지 못해 여러 방법을 찾았지만 결국 답을 맞추지 못해 다른 사람의 풀이를 참고하며 푼 문제.

다른 사람 풀이

import sys
input=sys.stdin.readline

N=int(input())
sequence=list(map(int,input().split()))
stack=[0]
answer=[-1]*N

for i in range(1,N):
    while stack and sequence[stack[-1]]<sequence[i]:
        answer[stack.pop()]=sequence[i]
    stack.append(i)
print(*answer)

이 문제는 스택으로 분류되어 있어 스택으로 풀 방법을 찾았지만 방법을 알 수가 없어서 헤맸다. 문제 풀이를 봤음에도 불구하고 이해하지 못해서 시간이 꽤 걸렸다.

우선 stack을 값을 넣는 용도가 아닌 인덱스 값을 넣는 용도로 사용한다.

처음엔 0 (인덱스 값)을 stack에 넣어주고 이후의 값들을 비교하여 비교하려는 값보다 작은 값은 stack에 추가해주고 보다 큰 값은 answer 리스트의 해당 배열에 넣어준 뒤 stack에 있던 값을 pop 해주는 것이다. 이 과정을 순서대로 적으면

0을 stack에 추가해준다 (첫 인덱스 값)
인덱스 1부터 N까지 차례대로 스택에 있는 값과 비교해준다.

1.1 stack의 제일 위에 있는 값을 인덱스로 한 배열값과 배열의 i번째 값과 비교한다. (stack에는 인덱스를 값으로 넣어준다.)
1.2 배열의 i번째 값이 더 크다면 answer의 stack[-1]를 인덱스로 한 위치에 배열의 i번째 값을 넣어주고
1.3 stack에 있는 값을 인덱스로 했을때 i번째 값보다 작은 값들을 차례로 pop해준다. (이때 i번째 값보다 작은 stack 값들은 1.2의 로직을 수행해야한다.)
1.4 배열의 i번째 값이 더 작다면 stack에 i를 추가해준다. (stack에는 인덱스가 들어가야하니까!!!)
모든 반복이 끝나면 answer값을 출력해준다.
(배열을 출력할 때, print에 *를 넣어주면 공백을 추가해서 출력해준다!)