백준 17298번 오큰수

https://www.acmicpc.net/problem/17298

17298번: 오큰수

 

 

문제 설명

길이가 $N$인 수열 A가 주어집니다.

A의 i번째 원소 $A_i$의 오큰수는 자신보다 오른쪽에 있고 큰 수 중 가장 가까운 수입니다. 

수열의 모든 원소의 오큰수를 찾아 출력하는 문제입니다.

 

풀이

각 원소의 오큰수를 찾을 때마다 오른쪽 원소를 모두 탐색하면 $O(N^2)$로 시간초과일 것입니다.

스택을 이용하면 $O(N)$만에 모든 원소의 오큰수를 찾을 수 있습니다.

 

수열 A의 오른쪽에서부터 원소를 탐색하며

1. 스택이 비어있거나 스택의 최상단 원소가 자신보다 클 때까지 pop연산을 수행합니다.

2. 스택이 비어있다면 현재 원소의 오큰수는 -1이며, 아니라면 스택의 최상단 원소가 오큰수입니다.

3. 현재 원소를 스택에 삽입합니다.

 

현재 원소를 $A_i$, 현재 원소보다 오른쪽에 있고 작거나 같은 원소를 $A_j$라고 하겠습니다.

$A_j$는 절대로 $A_i$보다 왼쪽에 있는 원소의 오큰수가 될 수 없습니다.

$A_i$보다 왼쪽의 원소가 $A_j$보다 작다면 $A_i$보다도 작으므로 더 가까운 $A_i$가 오큰수가 되기 때문입니다.

따라서 스택이 내림차순으로 유지되도록 위와 같이 관리하면 정답을 찾을 수 있습니다.

 

코드

#include <iostream>
#include <algorithm>
#include <stack>
#define S 1000005
using namespace std;
int arr[S], ans[S];
stack<int> stk;
int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL); cout.tie(NULL);
    int n;
    cin >> n;
    for (int i = 0; i < n; i++)
        cin >> arr[i];
    for (int i = n - 1; i >= 0; i--) {
        while (!stk.empty() and stk.top() <= arr[i])
            stk.pop();
        if (stk.empty())
            ans[i] = -1;
        else
            ans[i] = stk.top();
        stk.push(arr[i]);
    }
    for (int i = 0; i < n; i++)
        cout << ans[i] << ' ';
    return 0;
}