[알고리즘 / C++] 프로그래머스 - 뒤에 있는 큰 수 찾기

오늘은 프로그래머스의 뒤에 있는 큰 수 찾기 문제를 풀어보았다.

이 문제를 풀면서 처음으로 단조 스택(Monotonic Stack) 과
NGE(Next Greater Element) 유형을 알게 됐다.

처음에는 각 원소마다 오른쪽을 직접 탐색하면 될 것 같았지만,
그렇게 하면 시간 복잡도가 O(N^2)라 시간 초과로 실패하기 때문에 다른 풀이가 필요하다.

이 문제의 핵심은 스택으로 필요한 후보만 관리해서 O(N)에 해결하는 것이었다.

나는 이 문제를 두 가지 방식으로 풀어보았다.

• 앞에서부터 탐색 + 인덱스 스택
• 뒤에서부터 탐색 + 값 스택

---

NGE란?

NGE는 Next Greater Element의 줄임말로,
각 위치에서 오른쪽에 있으면서 처음 만나는 더 큰 수를 의미한다.

따라서, 어떤 값의 오른쪽을 확인할 때 자기보다 큰 값이 처음 등장하는 순간 그 값이 NGE가 된다.

예를 들어 아래와 같은 숫자 배열이 있다고 가정해보자
만약 NGE가 존재하지 않는다면 -1을 기록한다.

[9, 1, 5, 3, 6, 2]

각 원소 기준으로 보면 다음과 같다.

• 9 → 오른쪽에 더 큰 수가 없으므로 -1
• 1 → 오른쪽에서 처음 만나는 더 큰 수는 5
• 5 → 오른쪽에서 처음 만나는 더 큰 수는 6
• 3 → 오른쪽에서 처음 만나는 더 큰 수는 6
• 6 → 더 큰 수가 없으므로 -1
• 2 → 더 큰 수가 없으므로 -1

그러므로 숫자 배열의 NGE는 아래와 같은 결과가 나온다.

[-1, 5, 6, 6, -1, -1]

---

풀이 1 : 앞에서부터 탐색 + 인덱스 스택

이 방식은 현재 값이 이전 원소들의 답이 되는지를 확인하는 방식이다.
스택에는 아직 답을 못 찾은 인덱스를 저장한다.

현재 값이 스택 top이 가리키는 값보다 크면
현재 값이 그 인덱스의 NGE가 된다.

즉, 현재 값이 이전 원소들의 NGE를 확정하는 구조라고 볼 수 있다.

핵심 원리

• 스택에는 아직 정답이 정해지지 않은 인덱스가 들어간다.
• 현재 값을 보면서, 이전 원소들 중 현재 값이 NGE가 되는 원소들의 답을 확정한다.
• 현재 값 하나로 여러 이전 원소를 한꺼번에 처리한다.

앞에서부터 탐색하는 방식도 단조 스택인가?

처음 보면 이 방식은 스택에 인덱스를 저장하므로 단조 스택과는 조금 달라 보일 수 있다.

하지만 실제 비교는 인덱스 자체가 아니라
numbers[st.top()]처럼 그 인덱스가 가리키는 값 기준으로 이뤄진다.

스택 안에는 인덱스가 들어가지만
그 인덱스들이 가리키는 값들이 일정한 순서를 유지하도록 관리되고 있으므로
이 방식 역시 단조 스택 풀이로 볼 수 있다.

코드

vector<int> solution(vector<int> numbers) {
    // numbers 크기만큼 -1로 초기화
    vector<int> answer(numbers.size(), -1);

    // 인덱스를 넣는 스택
    stack<int> st;

    for (int i = 0; i < numbers.size(); ++i) {
        const int& num = numbers[i];

        // 스택이 안비었고
        // 스택 최상단 인덱스에 해당하는 숫자보다 현재 숫자가 크면 체크
        while (!st.empty() && numbers[st.top()] < num) {
            answer[st.top()] = num;
            st.pop();
        }
        // 검사한 인덱스 스택에 넣기
        st.push(i);
    }

    return answer;
}

---

풀이 2 : 뒤에서부터 탐색 + 값 스택

이 방식은 현재 원소의 답이 오른쪽 후보들 중 누군지 찾는 방식이다.
스택에는 뒤에 있는 큰 수 후보 값들을 저장한다.

현재 값보다 작거나 같은 값은 현재 원소의 NGE가 될 수 없으므로 제거한다.
정리 후 스택이 비어 있지 않다면 top이 현재 원소의 NGE가 된다.

즉, 현재 원소가 자기 답을 찾는 구조라고 볼 수 있다.

핵심 원리

• 스택에는 현재 위치의 오른쪽에 있는 값들 중, 큰 수 후보만 남긴다.
• 현재 값보다 작거나 같은 값은 정답 후보가 될 수 없으므로 제거한다.
• 정리 후 남아 있는 스택 top이 현재 원소의 정답이 된다.

코드

vector<int> solution(vector<int> numbers) {
    vector<int> answers(numbers.size(), -1);
    stack<int> st;

    // 뒤에서부터 체크
    // 특정 위치의 수가 stack에 있는 값보다 작으면 뒷 큰수로 저장
    for (int i = numbers.size() - 1; i >= 0; --i) {
        const int& num = numbers[i];

        // 만약 현재 탐색값이 stack의 top보다 크거나 같으면 pop
        while (!st.empty() && st.top() <= num) {
            st.pop();
        }

        // 위에서 정답후보만 남겼으면 top이 NGE
        if (!st.empty()) {
            answers[i] = st.top();
        }

        st.push(num);
    }

    return answers;
}

---

두 방식의 차이

앞에서부터 탐색

• 현재 값이 이전 원소들의 답을 정해준다.
• 스택에는 아직 답을 못 찾은 인덱스가 들어간다.

뒤에서부터 탐색

• 현재 원소가 자기 답을 스택에서 찾는다.
• 스택에는 뒤에 있는 큰 수 후보 값이 들어간다.

한 줄로 정리하면 다음과 같다.

• 앞에서부터 : 현재 값이 남의 답이 된다.
• 뒤에서부터 : 현재 원소가 자기 답을 찾는다.

---

마무리

이번 문제를 통해 단조 스택이 필요한 후보만 유지하는 방식이라는 걸 배웠다.

같은 문제도 앞에서부터 보느냐, 뒤에서부터 보느냐에 따라
풀이 관점이 달라질 수 있다는 점도 흥미로웠다.

또 앞에서부터 탐색하는 방식도 스택에 인덱스를 저장할 뿐,
실제로는 값의 순서를 관리한다는 점에서 단조 스택으로 볼 수 있다는 점도 새롭게 이해하게 됐다.

앞으로 비슷한 문제를 보면 단조 스택으로 풀 수 있는지 먼저 떠올려봐야겠다.