[LeetCode] 1381. Design a Stack With Increment Operation

https://leetcode.com/problems/design-a-stack-with-increment-operation/description/?envType=daily-question&envId=2024-09-30

leetcode - Design a Stack With Increment Operation

문제 유형 : 구현, 스택

문제 난이도 : Medium

 

문제

Design a stack that supports increment operations on its elements.

Implement the CustomStack class:

• CustomStack(int maxSize) Initializes the object with maxSize which is the maximum number of elements in the stack.
• void push(int x) Adds x to the top of the stack if the stack has not reached the maxSize.
• int pop() Pops and returns the top of the stack or -1 if the stack is empty.
• void inc(int k, int val) Increments the bottom k elements of the stack by val. If there are less than k elements in the stack, increment all the elements in the stack.

요소에 대해 증가 연산을 할 수 있는 스택을 디자인하시오.

최대 maxSize의 크기를 가집니다.

void push(int x): 현재 크기가 maxSize이내라면, x값을 톱에 추가

int pop(): 톱의 값을 pop하여 반환. 비어있으면 -1반환.

void inc(int k, int val): 바텀부터 k개의 요소들에 val만큼 값 증가. 요소가 k개 이하라면 전부 다 증가시키기.

 

풀이

간단하게 구현할 수 있겠지만, inc()함수를 O(1)에 푸는 방법이 있다.

inc배열을 만들어서, 증가값을 저장시키는 것이다.

Lazy Propagation이라는 기법이다.

 

push는 그냥 일반적인 배열을 이용한 스택 구현을 생각하면 되는데,

increment(k, val)함수에서 inc배열을 만들어서, 현재 스택의 톱이 0 이상이면, 즉 값이 하나라도 있다면, 스택의 bottom으로부터 k번째 자리 또는 스택의 톱, 즉 배열의 인덱스 min(k-1, top)에 증가값을 저장해놓는다.

 

그리고 pop연산에서 pop[top] + inc[top]의 값이 실제 반환되는 값이 된다.

앞의 인덱스들도 같은 increment가 되었기 때문에, inc[top-1]에 inc[top]을 누적시킨다.

그리고, inc[top]은 0으로 돌려놓고, top--로 톱의 위치를 조정한다.

그리고, 실제 값 값을 반환하면 된다.

 

 

코드

C++

#pragma GCC optimize("03", "unroll-loops");
static const int __ = [](){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    return 0;
}();

class CustomStack {
private:
    vector<int> arr;
    vector<int> inc;
    int top;
    int maxSize;
public:
    CustomStack(int maxSize) {
        top = -1;
        this -> maxSize = maxSize;
        arr.resize(maxSize);
        inc.resize(maxSize, 0);
    }
    
    void push(int x) {
        if(top == maxSize-1) return;
        top++;
        arr[top] = x;
    }
    
    int pop() {
        if(top < 0) return -1;
        int popped = arr[top] + inc[top];

        if(top > 0) {
            inc[top-1] += inc[top];
        }
        inc[top] = 0;
        top--;
        return popped;
    }
    
    void increment(int k, int val) {
        if(top >= 0) {
            int limit = min(k-1, top);
            inc[limit] += val;
        }
    }
};

/**
 * Your CustomStack object will be instantiated and called as such:
 * CustomStack* obj = new CustomStack(maxSize);
 * obj->push(x);
 * int param_2 = obj->pop();
 * obj->increment(k,val);
 */