프로그래머스 154539 - 뒤에 있는 큰 수 찾기 (Lv.5)

출처: https://school.programmers.co.kr/learn/courses/30/lessons/154539

// 프로그래머스 154539 - 뒤에 있는 큰 수 찾기 (Lv.5)
// https://school.programmers.co.kr/learn/courses/30/lessons/154539

// 문제 설명
// 배열의 각 원소에 대해, 자신보다 뒤에 있으면서 자신보다 크고 가장 가까운 수를
// "뒷 큰수"라 한다. 각 원소의 뒷 큰수를 담은 배열을 반환하라. 없으면 -1.

// 제약 조건
// 4 <= numbers 길이 <= 1,000,000
// 1 <= numbers[i] <= 1,000,000

// Example
// Input : [2, 3, 3, 5]
// Output: [3, 5, 5, -1]        (인덱스 1의 3은 뒤의 3이 '크지 않아' 5까지 기다린다)
//
// Input : [9, 1, 5, 3, 6, 2]
// Output: [-1, 5, 6, 6, -1, -1]

// 접근 — 단조 감소 스택에 "답을 기다리는 인덱스"를 쌓는다
// 이중 루프는 최악 10^6 * 10^6 = 10^12로 즉사한다. 관찰 하나로 O(n)이 된다:
// 앞에 있는 작은 수는 자기보다 큰 수가 나오는 순간 전부 같은 값으로 답이 확정된다.
// 1) 왼쪽부터 순회하며 "아직 답을 못 찾은 인덱스"를 스택에 쌓는다.
//    스택 안의 값은 항상 내려가는 순서(단조 감소)가 유지된다.
// 2) 현재 값이 스택 top의 값보다 크면 그 인덱스의 답은 현재 값 — 꺼내며 확정한다.
//    같은 값일 때는 꺼내지 않는다("크면서"이므로 3 뒤의 3은 답이 아니다).
// 3) 끝까지 남은 인덱스는 뒷 큰수가 없으므로 -1(초기값 그대로).
// 각 인덱스는 한 번 push, 최대 한 번 pop → 시간 O(n), 공간 O(n)

#include <vector>
#include <stack>

using namespace std;

vector<int> solution(vector<int> numbers) {
    int n = numbers.size();
    vector<int> answer(n, -1);                 // 못 찾으면 -1이 그대로 답
    stack<int> st;                             // 뒷 큰수를 기다리는 인덱스들 (값 기준 단조 감소)

    for (int i = 0; i < n; i++) {
        // '<' 이어야 한다 — '<='로 쓰면 같은 값을 뒷 큰수로 잘못 확정한다
        while (!st.empty() && numbers[st.top()] < numbers[i]) {
            answer[st.top()] = numbers[i];
            st.pop();
        }
        st.push(i);
    }
    return answer;
}

정리

• 제약이 답을 지정하는 문제다. 길이 10^6이면 O(n^2) = 10^12로 이중 루프는 시작부터 탈락이고, 값 범위도 10^6이라 정렬해도 원래 위치 정보가 필요해 이득이 없다. 남는 건 한 번 훑으면서 O(n)에 끝내는 구조뿐 — 이 지점에서 스택이 나온다.
• 관점을 뒤집는 게 핵심이다. “각 원소가 자기 답을 찾아 뒤를 뒤진다”가 아니라 “현재 원소가 앞에서 기다리던 애들의 답을 확정해준다”. [9,1,5,...]에서 5를 만나면 1의 답이 확정되고, 6을 만나면 3과 5의 답이 한꺼번에 확정된다. 답을 확정해주는 방향으로 보면 각 원소를 한 번만 봐도 된다.
• 스택에는 값이 아니라 인덱스를 담는다. answer[st.top()] = ...으로 답을 써야 하니 위치가 필요하고, 값은 numbers[st.top()]으로 언제든 꺼낼 수 있다. 값만 담으면 어디에 써야 할지 모른다 — 뒤로 미룬 작업을 스택에 쌓을 때는 “무엇을” 이 아니라 “어디에” 를 쌓는다.
• 스택이 항상 단조 감소로 유지되는 이유를 알면 while이 안전하다는 걸 안다. push 직전에 자기보다 작은 것을 모두 꺼내므로, 남은 top은 반드시 새 값 이상이다 → 불변식이 스스로 유지된다. 그래서 while이 도는 동안 꺼내는 값들은 모두 numbers[i]보다 작고, 그중 가장 가까운 큰 수가 numbers[i]임이 보장된다(더 가까운 큰 수가 있었다면 그 시점에 이미 꺼내졌을 것이다).
• 부등호가 <인지 <=인지가 유일한 함정. 예제 1 [2,3,3,5]가 이걸 정확히 찍는다 — 인덱스 1의 3은 뒤의 3이 “크지 않으므로” 답이 아니고 5까지 기다려야 [3,5,5,-1]이 나온다. <=로 쓰면 [3,3,5,-1]이 되어 틀린다. 문제 문구 “자신보다 크면서“를 부등호에 그대로 옮기면 되는데, 등호를 무심코 붙이는 실수가 나오기 쉬운 자리다.
• 복잡도는 amortized로 봐야 한다. while이 안에 있어서 이중 루프처럼 보이지만, 각 인덱스는 평생 한 번 push되고 최대 한 번 pop되므로 while의 총 반복 횟수가 n을 넘지 않는다 → 전체 O(n). 최악 케이스 두 개가 이걸 반대로 보여준다: 내림차순 입력은 pop이 한 번도 안 일어나 스택에 10^6개가 다 쌓이고(그래서 공간 O(n)), 오름차순 입력은 매번 하나씩 꺼내며 스택 크기가 1로 유지된다. 둘 다 총 연산은 O(n)이다.
• -1을 answer의 초기값으로 깔아두면 “없으면 -1” 처리가 코드에서 사라진다. 루프가 끝난 뒤 스택에 남은 인덱스를 따로 순회해 -1을 채우는 후처리가 필요 없다 — 아무도 답을 확정해주지 않았다는 사실이 그대로 -1로 남는다.
• 검증: 예제 2개([3,5,5,-1], [-1,5,6,6,-1,-1])와 최악 케이스 2개(길이 10^6 내림차순 → 전부 -1, 길이 10^6 오름차순 → 마지막만 -1) 통과 (MSVC /std:c++17 컴파일·실행). 커리큘럼 98번.

핵심 요약 — “각자 뒤를 뒤진다”를 “현재 값이 앞에서 기다린 애들의 답을 확정해준다”로 뒤집으면, 답을 기다리는 인덱스를 단조 감소 스택에 쌓아 O(n)에 끝난다. 부등호는 <여야 하고(같은 값은 뒷 큰수가 아니다), answer를 -1로 초기화해두면 “없으면 -1” 후처리가 아예 사라진다.