Leetcode Daily Question 30/09/2024 - 1381. Design a Stack With Increment Operation

Problem Description:

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.

class CustomStack:
  def __init__(self, maxSize: int):
    self.maxSize = maxSize
    self.stack = []
    self.increment_stack = []
    
  
  def push(self, x: int):
    if len(self.stack) < self.maxSize:
      self.stack.appned(x)
      self.increment_stack.append(0)
        
    
  def pop(self):
    if not self.stack:
      return -1
    
    if len(self.increment_stack) >= 2:
           self.increment_stack[-2] += self.increment_stack[-1]
        
    return self.stack.pop() + self.increment_stack.pop()
  
  
  def increment(self, k: int, val: int):
    if not self.increment_stack:
      return -1
    
    index = min(len(self.increment_stack) - 1, k - 1)
    
    self.increment_stack[index] += val

Explanation:

The most straight forward way to do this is to keep track of a stack and add values to the indices less than k.

However, since we do not need to know every value of the stack at all times. We can keep another array to track of the top-most element affected by each increment operation.

This way, we do not need to iterate through the stack every time to increment the necessary values.

This is perhaps best illustrated by an example.

# Example
Consider the series of actions 
["CustomStack","push","push","pop","push","push","push","increment","increment","pop","pop","pop","pop"]
[[3],[1],[2],[],[2],[3],[4],[5,100],[2,100],[],[],[],[]]

CustomStack(3) -- creates a stack of max length 3

push(1) -- push 1 onto the stack, append 0 increment_stack to keep track of length and incremeneted value.
self.stack = [1] self.increment_stack = [0]

push(2) 
self.stack = [1,2] self.increment_stack = [0,0]

pop()
Since we are popping from stack we also need to pop from increment_stack
self.stack.pop() + self.increment_stack.pop() = 2 + 0 = 2

push(2)
self.stack = [1,2] self.increment_stack = [0,0]

push(3)
self.stack = [1,2,3] self.increment_stack = [0,0,0]

push(4)
This does nothing as we reached the max length allowed

increment(5,100)
Now since k > len(self.stack) all the elements needed to increment by 100.
This is the same as saying the last element in increment_stack needed to be incremented.
So, self.stack = [1,2,3] self.increment_stack = [0,0,100]

increment(2,100)
k < len(self.stack), so this affects the last two element in self.increment_stack only.
So we find the last elemented to be affected index = min(len(self.increment_stack) - 1, k - 1)
self.stack = [1,2,3] self.increment_stack = [0, 100, 100]

pop()
To calculate the final sum, we consider the sum of the last two elements. These values will tell us how much we
need to add to the popped value in stack.
self.increment_stack = [0, 100 + 100, 100] = [0, 200, 100]
self.stack.pop() + self.increment_stack.pop() = 103

pop()
self.increment_stack = [200, 200]
self.stack.pop() + self.increment_stack.pop() = 202

pop()
self.increment_stack = [200]
self.stack.pop() + self.increment_stack.pop() = 201

pop()
nothing in the stack, so returns -1.

Putting the discussing logic into code.

def __init__(self, maxSize: int):
    self.maxSize = maxSize
    self.stack = []
    self.increment_stack = []
    
  
def push(self, x: int):
  if len(self.stack) < self.maxSize:
    self.stack.appned(x)
    self.increment_stack.append(0)
        
    
def pop(self):
  if not self.stack:
    return -1
    
  if len(self.increment_stack) >= 2:
    self.increment_stack[-2] += self.increment_stack[-1]
        
  return self.stack.pop() + self.increment_stack.pop()
  
  
def increment(self, k: int, val: int):
  if not self.increment_stack:
    return -1
    
  index = min(len(self.increment_stack) - 1, k - 1)
    
  self.increment_stack[index] += val

The time complexity of every operation is O(1). Since we only keep track of the right most affected element by calling increment and calculate the necessary value only when needed.

The space complexity is O(n) as worst case is when stack is continuously filled to the max. increment_stack shadows stack, so also has O(n) space.

The full solution can be found near the top of the page.