C++ 증분 연산을 지원하는 커스텀 스택 설계하기

문제 개요

다음과 같은 연산을 지원하는 스택을 설계해 보겠습니다.

• CustomStack(int maxSize) — 스택에 저장할 수 있는 최대 원소 개수인 maxSize로 객체를 초기화합니다. 스택이 이미 maxSize에 도달했다면 더 이상 원소를 추가하지 않습니다.
• void push(int x) — 스택이 아직 maxSize에 도달하지 않았다면 x를 스택의 맨 위에 삽입합니다.
• int pop() — 스택의 맨 위 원소를 제거하고 그 값을 반환합니다. 스택이 비어 있으면 -1을 반환합니다.
• void inc(int k, int val) — 스택의 아래쪽(바닥부터) k개 원소를 val만큼 증가시킵니다. 스택에 있는 원소가 k개보다 적다면 전체 원소를 증가시킵니다.

접근 방법: 지연 갱신(Lazy Update) 기법

increment 연산이 호출될 때마다 하위 k개 원소를 일일이 갱신하면 최악의 경우 O(N) 시간이 걸립니다. 대신 별도의 증분 배열 inc를 두어 증가분을 기록해 두고, pop이 발생하는 순간 한 번에 반영하면 모든 연산을 O(1)에 처리할 수 있습니다.

풀이 절차는 다음과 같습니다.

• 배열 st(스택 본체)와 inc(증분 누적용 배열), 정수형 변수 cap을 선언합니다.
• 생성자에서 cap을 N으로 설정하고, inc를 크기 N + 10의 배열로 초기화합니다.
• push(x): 스택의 크기가 cap이 아니면 x를 st에 추가합니다.
• pop()은 다음과 같이 동작합니다.
  • st가 비어 있으면 -1을 반환합니다.
  • 그렇지 않으면:
    • top 값에 inc[top 인덱스]를 더해 누적된 증분을 실제 값에 반영합니다.
    • 스택에 원소가 하나 이상 남아 있다면 inc[size - 2]에 inc[size - 1]을 더해 증분을 바로 아래 원소로 전파합니다.
    • inc[size - 1]을 0으로 초기화합니다.
    • x := st의 마지막 원소를 꺼내 반환합니다.
• inc(k, val)는 다음과 같이 동작합니다.
  • k를 1 감소시켜 0 기반 인덱스로 변환합니다.
  • k := min(k, st.size() - 1)로 조정합니다.
  • k < 0이면(스택이 비어 있으면) 그대로 반환합니다.
  • inc[k]에 val을 더합니다.

동작 원리

핵심 아이디어는 "증분을 해당 위치에 기록해 두고, pop이 발생하는 순간에 실제 값에 반영하는 것"입니다. inc[i]에는 i번째 원소부터 그 위의 모든 원소에 적용되어야 할 증분이 누적됩니다. 따라서 pop 시 top에 inc[top]을 더해 주면 되고, 원소가 제거된 후에는 그 증분을 바로 아래 원소의 inc 값에 이어받게 하여 아래쪽 원소들도 동일한 증분의 영향을 받도록 처리합니다.

C++ 구현 예제

다음 구현을 통해 더 잘 이해해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;

class CustomStack {
public:
    vector<int> st;
    vector<int> inc;
    int cap;

    CustomStack(int N) {
        cap = N;
        inc = vector<int>(N + 10);
    }

    void push(int x) {
        if(st.size() == cap) return;
        st.push_back(x);
    }

    int pop() {
        if(st.empty()) return -1;
        else {
            st.back() += inc[st.size() - 1];
            if(st.size() - 1 > 0) {
                inc[st.size() - 2] += inc[st.size() - 1];
            }
            inc[st.size() - 1] = 0;
            int x = st.back();
            st.pop_back();
            return x;
        }
    }

    void increment(int k, int val) {
        k--;
        k = min(k, (int)st.size() - 1);
        if(k < 0) return;
        inc[k] += val;
    }
};

int main(){
    CustomStack ob(3);
    ob.push(1);
    ob.push(2);
    cout << ob.pop() << endl;
    ob.push(2);
    ob.push(3);
    ob.push(4);
    ob.increment(5, 100);
    ob.increment(2, 100);
    cout << ob.pop() << endl;
    cout << ob.pop() << endl;
    cout << ob.pop() << endl;
    cout << ob.pop() << endl;
    return 0;
}

출력 결과

2
103
202
201
-1

결과 해석

처음 pop()은 맨 위의 2를 반환합니다. 이후 push(2), push(3)으로 스택이 [1, 2, 3]이 되고, 용량이 3이므로 push(4)는 무시됩니다. increment(5, 100)은 원소가 5개보다 적으므로 전체 원소에 100을 더하고, increment(2, 100)은 아래 2개 원소에 100을 더해 스택은 [201, 202, 103]이 됩니다. 따라서 pop()은 차례로 103, 202, 201을 반환하고, 마지막 pop()은 빈 스택에서 호출되어 -1을 반환합니다.

복잡도 분석

push, pop, increment 세 연산 모두 O(1) 시간에 수행됩니다. 공간 복잡도는 스택 본체와 증분 배열을 위해 O(N)입니다.