Taro Logo

Design a Stack With Increment Operation

Medium
Google logo
Google
2 views
Topics:
Stacks

Design a stack data structure with the following operations:

  1. CustomStack(int maxSize): Initializes the stack with a maximum size of maxSize. The stack should not be able to hold more than maxSize elements.
  2. void push(int x): Pushes the element x onto the top of the stack. If the stack is full (i.e., already contains maxSize elements), do nothing.
  3. int pop(): Removes and returns the element at the top of the stack. If the stack is empty, return -1.
  4. void increment(int k, int val): Increments the bottom k elements of the stack by val. If there are fewer than k elements in the stack, increment all the elements in the stack.

Example:

CustomStack stk = new CustomStack(3); // Stack is Empty []
stk.push(1);                          // stack becomes [1]
stk.push(2);                          // stack becomes [1, 2]
stk.pop();                            // return 2, stack becomes [1]
stk.push(2);                          // stack becomes [1, 2]
stk.push(3);                          // stack becomes [1, 2, 3]
stk.push(4);                          // stack remains [1, 2, 3] (maxSize reached)
stk.increment(5, 100);                // stack becomes [101, 102, 103]
stk.increment(2, 100);                // stack becomes [201, 202, 103]
stk.pop();                            // return 103, stack becomes [201, 202]
stk.pop();                            // return 202, stack becomes [201]
stk.pop();                            // return 201, stack becomes []
stk.pop();                            // return -1, stack is empty

Explain how you would implement this CustomStack class with efficient time complexity for all operations. Consider edge cases and constraints while designing your solution.

Solution


Clarifying Questions

When you get asked this question in a real-life environment, it will often be ambiguous (especially at FAANG). Make sure to ask these questions in that case:

  1. What are the constraints on the size of `maxSize`, and the values of `val` and `k` in the `increment` operation?
  2. What should happen if `k` is larger than the current stack size? Should I increment all elements in the stack or just the bottom `k`?
  3. Will there be multiple calls to `push`, `pop`, and `increment` after the stack reaches its maximum size?
  4. Is `maxSize` guaranteed to be a non-negative integer?
  5. Are we optimizing for space or time complexity, given the potential trade-offs between using a lazy propagation approach versus eagerly updating the stack during the increment operation?

Brute Force Solution

Approach

We need to create a special stack that can also increase the values of existing items. A brute force method would directly implement each operation as simply as possible, without any clever tricks or optimizations.

Here's how the algorithm would work step-by-step:

  1. To push a new number onto the stack, just add it to the top of the stack.
  2. To remove the top number, simply take it off the top of the stack.
  3. To increment the bottom numbers, go through each of those numbers from the bottom up and add the given value to them. If there are fewer numbers in the stack than what needs to be incremented, then just increment the ones that exist.

Code Implementation

class CustomStack:

    def __init__(self, max_size):
        self.stack = []
        self.max_size = max_size

    def push(self, value):
        if len(self.stack) < self.max_size:
            self.stack.append(value)

    def pop(self):
        if self.stack:
            return self.stack.pop()
        else:
            return -1

    def increment(self, number_of_elements, value):
        # Determine how many elements to increment
        increment_range = min(number_of_elements, len(self.stack))

        # Iterate through the bottom elements of the stack
        for index in range(increment_range):
            # Increment each element within the range.
            self.stack[index] += value

Big(O) Analysis

Time Complexity
O(k)The push and pop operations take constant time, O(1), as they simply involve adding or removing an element from the top of the stack. The increment operation iterates through the bottom k elements of the stack, where k is the number of elements to increment. Therefore, the increment operation takes O(k) time, where k is a parameter of the increment function, independent of the total number of elements 'n' in the stack. If k can grow proportionally to n, then increment will be O(n). Otherwise, the complexity depends on the value of k.
Space Complexity
O(1)The brute force method described operates directly on the stack's underlying data structure. The push and pop operations don't require any auxiliary space. The increment operation, while iterating through stack elements, modifies them in place and doesn't create any new data structures to store intermediate results or temporary copies. Therefore, the space complexity is constant, independent of the number of elements N in the stack.

Optimal Solution

Approach

We need to design a special stack that can not only push and pop values, but also increment several of its bottom elements. The clever idea is to avoid changing many stack values individually when incrementing; instead, we store increment amounts separately for efficiency.

Here's how the algorithm would work step-by-step:

  1. First, create two containers. One will hold the stack's values, like a normal stack. The other will hold the increment values corresponding to each element in the stack.
  2. When pushing a new value onto the stack, also push a 'zero' into the increment container. This means, initially, the new value doesn't need to be incremented by anything.
  3. When popping a value, retrieve both the stack value and its corresponding increment value. Add the increment to the stack value and return the result. Also, if there are any values underneath in the stack, add current increment value to the item underneath the top item.
  4. For the increment operation, take the number of elements to increment and the increment value. Add the increment value to the corresponding element in the increment container. Elements beyond the specified number of elements are not affected.
  5. By storing increments separately, we can apply increments to many elements quickly, without needing to modify each individual stack value right away.

Code Implementation

class CustomStack:

    def __init__(self, max_size: int):
        self.max_size = max_size
        self.stack_values = []
        self.increment_values = []

    def push(self, value: int) -> None:
        if len(self.stack_values) < self.max_size:
            self.stack_values.append(value)
            # Initialize increment value to 0 for new element
            self.increment_values.append(0)

    def pop(self) -> int:
        if not self.stack_values:
            return -1

        stack_index = len(self.stack_values) - 1
        value_to_pop = self.stack_values.pop()
        increment_to_pop = self.increment_values.pop()

        value_to_return = value_to_pop + increment_to_pop

        # Propagate increment to the element below.
        if self.increment_values:
            self.increment_values[-1] += increment_to_pop

        return value_to_return

    def increment(self, number_of_elements: int, value: int) -> None:
        # Increment only the bottom 'number_of_elements' elements
        number_of_elements_to_inc = min(number_of_elements, len(self.stack_values))

        if number_of_elements_to_inc > 0:
            # Apply increment to the appropriate element.
            self.increment_values[number_of_elements_to_inc - 1] += value

Big(O) Analysis

Time Complexity
O(k)The push operation involves appending to both the stack and the increments array, which takes O(1) time. The pop operation also involves popping from both the stack and the increments array. Additionally, it adds the top increment value to the element beneath it, which takes O(1). The increment operation adds a value to the increments array at a specific index, up to k elements from the bottom of the stack, where k is the number of elements to increment. Therefore, the increment operation takes O(k) time, as it iterates through a portion of the array to modify the corresponding increment values.
Space Complexity
O(N)The solution uses two containers to store data. One container stores the stack's values, and the other stores the increment values corresponding to each element in the stack. Both containers grow linearly with the number of elements pushed onto the stack. Therefore, if N is the maximum number of elements ever present in the stack, the auxiliary space used is proportional to N, resulting in a space complexity of O(N).

Edge Cases

CaseHow to Handle
maxSize is zeroThe CustomStack should be initialized but `push` operations will have no effect.
Pushing more elements than maxSize allowsOnly the first maxSize elements should be stored; subsequent `push` calls should be ignored.
Popping from an empty stack`pop` should return -1 if the stack is empty.
Incrementing with k = 0The increment operation should have no effect as no elements are incremented.
Incrementing with k larger than the stack sizeAll elements in the stack should be incremented by val.
Large val during increment causing integer overflowUse a data type (e.g., long) that can accommodate larger numbers or perform overflow checks and handle appropriately (e.g. throwing exception or clamping).
A sequence of pushes, pops, and increments that leave the stack empty then followed by an incrementIncrement should have no effect if the stack is empty after the pop sequence.
maxSize is a very large number (close to the integer limit)Ensure memory allocation doesn't exceed available resources, potentially leading to program crashes due to OOM or other issues; it might be necessary to consider alternative data structures or approaches if scalability is crucial.