백준 17928번: 오큰수 (C++)

문제

크기가 N인 수열 \(A\) = \(A_1\), \(A_2\), …, \(A_N\)이 있다. 수열의 각 원소 \(A_i\)에 대해서 오큰수 NGE(i)를 구하려고 한다. \(A_i\)의 오큰수는 오른쪽에 있으면서 \(A_i\)보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다. 그러한 수가 없는 경우에 오큰수는 -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\)의 원소 \(A_1\), \(A_2\), …, \(A_N\) (1 ≤ \(A_i\) ≤ 1,000,000)이 주어진다.
• 출력
  • 총 N개의 수 NGE(1), NGE(2), …, NGE(N)을 공백으로 구분해 출력한다.

코드

#include <iostream>
#include <stack>
using namespace std;

int main(void){
    ios::sync_with_stdio(0);
    cin.tie(0);

    int num;
    cin >> num;

    stack<pair<int, int> > s;
    int arr[1000001] = {0};
    for(int i = 1; i <= num; i++){
        int n;
        cin >> n;

        while(!s.empty() && s.top().second < n){
            arr[s.top().first] = n;
            s.pop();
        }
        s.push(make_pair(i, n));       
    }

    for (int i = 1; i <= num; i++){
        if (arr[i] == 0) arr[i] = -1;
        cout << arr[i] << " ";
    }
}

접근 및 풀이

이번 문제는 탑 문제와 비슷하지만 약간 다른 문제이다. 하지만 문제를 찬찬히 생각해보면 탑 문제와 접근하는 방법 자체는 다를 게 없다는 생각이 들 것이다.

먼저 나는 오큰수를 기록할 배열 arr을 만들어 0으로 초기화 해주었다. 그리고 스택은 수열의 index와 원소를 pair로 저장하도록 하였다. 그 다음 수열에서 원소를 순서대로 꺼내어 스택에 top에 있는 원소와 비교한다. 오큰수는 오른쪽에 있으면서 해당 수보다 큰 수 중에 가장 왼쪽에 있는 수라고 하였다. 만약 스택에 새로 들어온 원소가 스택의 top에 있는 원소보다 작을 경우에는 오큰수가 될 수 없으니 스택에 push해주고, top에 있는 원소보다 클 경우에는 오큰수가 되니 해당 원소를 top에 있는 index번째 배열에 기록하고 pop해주는 과정을 반복해준다. 이 과정은 새로 들어온 원소가 스택의 top에 있는 원소보다 작아질 때까지 반복된다.

만약 스택에 아무것도 없다면 비교할 원소가 없기 때문에 바로 push 해주고, 오큰수가 없는 경우에는 배열에 오큰수가 기록되지 않아 0으로 남아있을 것이기 때문에 문제의 조건에 따라 -1로 바꾼 뒤 arr 배열을 모두 출력해 주면 된다.