class CustomStack:
def __init__(self, maxSize: int):
self.stack = []
self.inc = []
self.maxSize = maxSize
def push(self, x: int) -> None:
if len(self.stack) < self.maxSize:
self.stack.append(x)
self.inc.append(0)
def pop(self) -> int:
if not self.stack:
return -1
idx = len(self.stack) - 1
res = self.stack.pop() + self.inc.pop()
if self.inc and idx > 0:
self.inc[idx - 1] += self.inc[idx]
return res
def increment(self, k: int, val: int) -> None:
n = min(k, len(self.stack))
if n > 0:
self.inc[n - 1] += val
class CustomStack {
public:
vector<int> stack, inc;
int maxSize;
CustomStack(int maxSize) {
this->maxSize = maxSize;
}
void push(int x) {
if (stack.size() < maxSize) {
stack.push_back(x);
inc.push_back(0);
}
}
int pop() {
if (stack.empty()) return -1;
int idx = stack.size() - 1;
int res = stack.back() + inc.back();
stack.pop_back();
int add = inc.back();
inc.pop_back();
if (!inc.empty()) inc.back() += add;
return res;
}
void increment(int k, int val) {
int n = min(k, (int)stack.size());
if (n > 0) inc[n - 1] += val;
}
};
class CustomStack {
private int[] stack;
private int[] inc;
private int size, maxSize;
public CustomStack(int maxSize) {
this.maxSize = maxSize;
this.stack = new int[maxSize];
this.inc = new int[maxSize];
this.size = 0;
}
public void push(int x) {
if (size < maxSize) {
stack[size] = x;
inc[size] = 0;
size++;
}
}
public int pop() {
if (size == 0) return -1;
int idx = size - 1;
int res = stack[idx] + inc[idx];
if (idx > 0) inc[idx - 1] += inc[idx];
size--;
return res;
}
public void increment(int k, int val) {
int n = Math.min(k, size);
if (n > 0) inc[n - 1] += val;
}
}
var CustomStack = function(maxSize) {
this.stack = [];
this.inc = [];
this.maxSize = maxSize;
};
CustomStack.prototype.push = function(x) {
if (this.stack.length < this.maxSize) {
this.stack.push(x);
this.inc.push(0);
}
};
CustomStack.prototype.pop = function() {
if (this.stack.length === 0) return -1;
let idx = this.stack.length - 1;
let res = this.stack.pop() + this.inc.pop();
if (this.inc.length > 0 && idx > 0) {
this.inc[idx - 1] += this.inc[idx];
}
return res;
};
CustomStack.prototype.increment = function(k, val) {
let n = Math.min(k, this.stack.length);
if (n > 0) this.inc[n - 1] += val;
};
You are asked to design a stack-like data structure with a twist. The stack should support three operations:
push(x)
: Add element x
to the top of the stack if it is not full.pop()
: Remove and return the element on the top of the stack, or return -1
if the stack is empty.increment(k, val)
: Increment the bottom k
elements of the stack by val
. If there are fewer than k
elements, increment all of them.maxSize
, and cannot exceed this capacity. Each operation should be efficient, especially increment
, which should not take linear time per operation.
At first glance, implementing push
and pop
is straightforward using a standard stack or array. The main challenge is the increment
operation. If we naively loop through the first k
elements and add val
to each, it would take O(k) time per operation, which could be slow if k
is large and called frequently.
We need to optimize increment
so it works efficiently, ideally in constant time. The key insight is to avoid updating the stack elements directly during increment
. Instead, we can keep track of the increments separately and only apply them when needed (i.e., during pop
).
This is similar to the "lazy propagation" technique used in segment trees, where updates are deferred until the value is actually needed. By being clever about when and how we apply increments, we can ensure all operations are efficient.
Let's break down the efficient solution:
inc
) to track the pending increment for each position.inc
array (since there's no increment for the new element yet).inc
to the popped element.inc[top]
to inc[top-1]
).val
to inc[k-1]
. This means that when any of the bottom k
elements are popped, they will get the increment at that time.This approach ensures all operations run in O(1) time, as we avoid looping through the stack for increments.
Let's walk through an example with maxSize = 3
:
push(1)
: stack = [1], inc = [0]push(2)
: stack = [1, 2], inc = [0, 0]increment(2, 100)
: inc = [0, 100] (now, bottom 2 elements will get +100 when popped)push(3)
: stack = [1, 2, 3], inc = [0, 100, 0]increment(3, 100)
: inc = [0, 100, 100]pop()
: pop 3, add inc[2] (100), result = 3 + 100 = 103. Carry inc[2] to inc[1]: inc = [0, 200]pop()
: pop 2, add inc[1] (200), result = 2 + 200 = 202. Carry inc[1] to inc[0]: inc = [200]pop()
: pop 1, add inc[0] (200), result = 1 + 200 = 201Notice how increments are efficiently accumulated and applied only when needed.
Brute-force Approach:
push
and pop
: O(1)increment
: O(k) per operation (could be O(n) if k ≈ n)increment
is called often with large k
.inc
array):
push
, pop
, increment
): O(1) timemaxSize
, for the stack and increment arrays.
The optimized approach is much faster, especially for frequent increment
operations.
The problem of designing a stack with an efficient increment operation is elegantly solved by using an auxiliary array to track pending increments. By deferring the increment operation and only applying it during pop
, we achieve constant time for all operations. This approach avoids unnecessary looping and demonstrates how clever bookkeeping can significantly optimize seemingly expensive operations.