【基础篇 - Day 3】 2020-11-03 - 1381. 设计一个支持增量操作的栈（01. 数组，栈，队列 ）

[patrickz93]

思路

用数组模拟栈，用一个变量记录数组末尾元素的下标，实现入栈、出栈。 使用另外一个数组用来记录栈底n个元素的增加值情况，元素出栈时，两个数组相应位置上的元素清零，末尾元素下标做减法。 只有当元素出栈的时候，获取该下标所有的incr结果，达到O（1）。

代码

class CustomStack {

    int[] arr;
    int[] incrArr;
    int currentIndex = -1;
    int maxSize;

    public CustomStack(int maxSize) {
        arr = new int[maxSize];
        incrArr = new int[maxSize];
        this.maxSize = maxSize;
    }

    public void push(int x) {
        if(currentIndex < maxSize - 1) {
            arr[++currentIndex] = x;
        }
    }

    public int pop() {
        if(currentIndex == -1) {
            return -1;
        }

        int val = arr[currentIndex] + incrArr[currentIndex];
        incrArr[currentIndex] = 0;
        arr[currentIndex] = 0;
        currentIndex--;
        return val;
    }

    public void increment(int k, int val) {
        if(currentIndex > -1) {
            int length = k-1 >= currentIndex ? currentIndex : k-1;
            for (int i = 0; i <= length; i++) {
                incrArr[i] += val;
            }
        }
    }
}

复杂度分析

• 时间复杂度：O(1)，除increment操作最坏情况下，均是O(1)，因为是通过下标直接访问数组元素；incr最坏情况下是O(n)，其中 N 为数组长度
• 空间复杂度：O(N)，其中 N 为数组长度。

[patrickz93]

increment优化至O(1)

public int pop() {
        if(currentIndex == -1) {
            return -1;
        }

        int val = arr[currentIndex] + incrArr[currentIndex];
        if (currentIndex > 0) {
            incrArr[currentIndex - 1] += incrArr[currentIndex];
        }
        incrArr[currentIndex] = 0;
        currentIndex--;
        return val;
    }

    public void increment(int k, int val) {
        int incrArrLastIndex = k-1 >= currentIndex ? currentIndex : k-1;
        if(incrArrLastIndex > -1) {
            incrArr[incrArrLastIndex] += val;
        }
    }