[BOJ] 17298번: 오큰수

https://www.acmicpc.net/problem/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이다.

 

풀이

샤라웃투@

N = int(input())
A = list(map(int, input().split()))
stack = [A[-1]]
ans = [-1]

for i in range(N - 2, -1, -1):
    while stack:
        if A[i] < stack[-1]:
            ans.append(stack[-1])
            break
        else: stack.pop()
    else:
        ans.append(-1)
    stack.append(A[i])

print(' '.join(map(str, ans[::-1])))

1. N과 A를 입력받는다.

2. A의 마지막 요소를 가지는 스택과 -1을 가지는 정답 리스트를 생성한다.

3. 뒤에서 두번째 수부터 0까지 for문으로 반복한다.

4. stack이 비어있지 않을 때, 만약 A[i]보다 스택의 가장 끝 값이 크다면, "오큰수"에 부합하므로 정답 리스트인 ans에 해당 값을 추가한다. 만약 그렇지 않다면 stack을 pop하고 다시 비교한다.

5. 만약 오큰수에 부합하는 수가 없으며 스택이 비어있다면 -1을 ans에 추가한다.

6. 현재 수 A[i]를 ans에 추가한다.

7. ans 배열을 뒤에서부터 출력한다.

 

for문의 i에 대응하는 A[i]를 현재 수라고 하면, stack에는 A[i+1]보다 큰 수와 A[i+1]만이 남아있게 되므로 오큰수를 적절히 찾을 수 있게 된다.