🔗 🟡 1381. Design a Stack With Increment Operation 1286

tags: `Weekly Contest 180` `陣列(Array)` `堆疊(Stack)` `設計(Design)`

題意

設計一個支持對其元素進行遞增操作的堆疊(Stack)。

實現 CustomStack 類別：

CustomStack(int maxSize) 初始化對象，maxSize 是堆疊中的最大元素數量。
void push(int x) 如果堆疊未達到 maxSize，將 x 添加到堆疊頂部。
int pop() 彈出並返回堆疊頂部的元素，如果堆疊為空，則返回 -1。
void inc(int k, int val) 將堆疊底部的 k 個元素遞增 val。如果堆疊中的元素少於 k 個，則遞增所有堆疊中的元素。

約束條件

$1 \leq \text{maxSize}, x, k \leq 1000$
$0 \leq \text{val} \leq 100$
最多會分別對 increment, push 和 pop 操作各調用 $1000$ 次。

思路：模擬(Simulation) + 懶標記(Lazy Tag)

除了 inc 操作，其他操作都是標準的堆疊操作，因此我們只需要關注 inc 操作。

首先注意到 inc 操作會影響到堆疊底部的元素，但標準的堆疊操作 push 和 pop 只會影響到頂部的元素。因此我們需要使堆疊中的所有元素都可以被直接訪問，這裡使用一個陣列 $st$ 來模擬堆疊，並使用一個指標 $top$ 來表示堆疊的頂部位置。

再來思考如何實現 inc 操作。最暴力的方式是每次 inc 操作都遍歷一次堆疊，將底部的 k 個元素遞增 val。但這樣會導致每次 inc 操作的時間複雜度變為 $O(n)$ ，其中 $n$ 是堆疊的最大容量。

為了降低時間複雜度，可以使用懶標記(Lazy Tag) 的思想，只在需要時才進行值的更新。具體來說，引入一個與堆疊等長的 $lazy$ 陣列，其中 $lazy[i]$ 表示 $st[0], st[1], \cdots, st[i]$ 需要增加的值。當執行 inc 操作時，我們只需將 $val$ 加到第 $k$ 個元素（即下標 $k-1$ ）的位置的 $lazy$ 標記上，而不需要立即更新所有底部 $k$ 個元素。

在 pop 操作時，當彈出頂部元素時，我們將其對應的 $lazy$ 值加到該元素上，並將該 $lazy$ 值傳遞給下一個元素（即 $top - 1$ 位置）。這樣，我們只在實際需要時才進行值的更新，從而減少不必要的操作。

此外，需要注意 inc 操作的邊界條件，當 $k$ 超過堆疊的元素數量時，我們只需要將 $val$ 加到堆疊頂部的元素上。

複雜度分析

時間複雜度：全為 $\mathcal{O}(1)$ $O (1)$
- push 操作： $\mathcal{O}(1)$
- pop 操作： $\mathcal{O}(1)$
- inc 操作： $\mathcal{O}(1)$
空間複雜度： $\mathcal{O}(maxSize)$ $O (ma x S i ze)$
- 需要 $O(maxSize)$ 的空間複雜度來保存堆疊和懶標記。

class CustomStack:

    def __init__(self, maxSize: int):
        self.top = -1
        self.st = [0] * maxSize
        self.lazy = [0] * maxSize # lazy tag

    def push(self, x: int) -> None:
        if self.top + 1 == len(self.st):
            return
        self.top += 1
        self.st[self.top] = x

    def pop(self) -> int:
        if self.top == -1:
            return -1
        res = self.st[self.top] + self.lazy[self.top]
        if self.top != 0: # 不是最底層的話，把 lazy tag 往下傳
            self.lazy[self.top - 1] += self.lazy[self.top]
        self.lazy[self.top] = 0 # 清掉 lazy tag
        self.top -= 1
        return res
        
    def increment(self, k: int, val: int) -> None:
        idx = min(k - 1, self.top) # 最多只會影響到下標為 top 的元素
        if idx >= 0:
            self.lazy[idx] += val

class CustomStack {
public:
    int top;
    vector<int> st;
    vector<int> lazy;
    CustomStack(int maxSize) {
        top = -1;
        st.resize(maxSize);
        lazy.resize(maxSize);
    }
    
    void push(int x) {
        if (top + 1 == st.size()) return;
        top++;
        st[top] = x;
    }
        
    int pop() {
        if (top == -1) return -1;
        int res = st[top] + lazy[top];
        if (top != 0) lazy[top - 1] += lazy[top]; // Propagate Lazy tag
        lazy[top] = 0; // Clear Lazy tag
        top--;
        return res;

    }
    
    void increment(int k, int val) {
        int idx = min(k - 1, top);
        if (idx >= 0) lazy[idx] += val;
    }
};