We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

SimplyLeet
TagsDisclaimersPrivacyTermsContactAbout Us

Design a Stack With Increment Operation

Difficulty: Medium

Problem Description

Design a stack that supports increment operations on its elements. Implement the CustomStack class which includes methods to push elements, pop elements, and increment the bottom k elements by a specified value.

Key Insights

  • The stack should support typical stack operations: push and pop.
  • An additional operation, increment, allows incrementing the bottom k elements.
  • We need to manage the maximum size of the stack and handle edge cases where k exceeds the number of elements in the stack.

Space and Time Complexity

Time Complexity:

  • push: O(1)
  • pop: O(1)
  • inc: O(k) in the worst case (where k is the number of elements to increment)

Space Complexity: O(n), where n is the maximum size of the stack.

Solution

We will use an array to represent the stack and an additional array to track increment values for each element. The push method will add elements to the stack if it hasn't reached its maximum size. The pop method will remove and return the top element, applying any increment from the corresponding index if necessary. The inc method will add the specified value to the increment array for the bottom k elements, allowing us to efficiently manage increments without immediately applying them to the stack.

Code Solutions

class CustomStack:
    def __init__(self, maxSize: int):
        self.maxSize = maxSize
        self.stack = []
        self.increment = [0] * maxSize  # Increment array

    def push(self, x: int) -> None:
        if len(self.stack) < self.maxSize:
            self.stack.append(x)

    def pop(self) -> int:
        if not self.stack:
            return -1
        index = len(self.stack) - 1
        value = self.stack.pop() + self.increment[index]  # Apply increment
        if index > 0:
            self.increment[index - 1] += self.increment[index]  # Pass increment to the next element
        self.increment[index] = 0  # Reset the increment for the popped element
        return value

    def inc(self, k: int, val: int) -> None:
        # Increment the first k elements
        limit = min(k, len(self.stack))
        if limit > 0:
            self.increment[limit - 1] += val  # Increment the last of the first k elements
← Back to All Questions