문제

크기가 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이다.

풀이

1. for문을 이용한 풀이 (시간초과)

A[i]보다 큰 것 중 가장 오른쪽에 있는 수를 구하는 것이므로
(1) i 기준으로 오른쪽에 있는 것 중
(2) A[i]보다 크기만 하면 된다
따라서 i+1 부터 시작해서 for문을 돌려서 A[i]보다 큰 것이 나오면 출력하고, break 하는 방식으로 접근했다.

하지만 이 방법은 시간 초과가 발생했다.

2. stack을 이용한 풀이

stack에 오큰수가 없는 경우의 인덱스를 저장하는 방식의 풀이.

(1) num을 입력받음
(2) - 1. stack이 비어있는 경우, 그냥 stack에 push
(2) - 2. 아닌 경우, stack의 top에 해당하는 배열 A의 인덱스와 비교
(2) - 2 - 1. top이 더 클 경우, 오큰수이므로 배열 A에 저장하고 pop
(2) - 2 - 2. top이 더 작을 경우, 오큰수가 아니므로 스택에 push
(3) stack에 남은 A의 인덱스들의 값을 -1로 바꿔줌
(4) 오큰수 출력

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

int A[1000001]; // 오큰수를 저장하는 배열
int main()
{
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    int N;
    cin >> N; // 배열의 크기 입력받기
    stack<int> st; // -1을 출력하는 경우의 배열 A의 인덱스를 저장하는 스택
    for (int i = 0; i < N; i++) {
        int num;
        cin >> num; // 수 입력받기
        A[i] = num;
        while (true) { // 스택이 비어있는 경우, push
            if (st.empty()) {
                st.push(i);
                break;
            }
            if (num > A[st.top()]) { // 오큰수일 경우 num을 넣어주고 pop
                A[st.top()] = num;
                st.pop();
            }
            else { // num이 작거나 같으면 push
                st.push(i);
                break;
            }
        }
    }
    while (!st.empty()) { // 스택에 남아 있는 index에는 -1을 넣어줌
        A[st.top()] = -1;
        st.pop();
    }
    for (int i = 0; i < N; i++) { // 오큰수 출력
        cout << A[i] << ' ';
    }
    return 0;
}

이 방식의 풀이는 O(N)의 시간복잡도로, 시간 초과가 발생하지 않는다.