Stack of plates

Problem

Imagine a (literal) stack of plates. If the stack gets too high, it might topple. Therefore, in real life, we would likely start a new stack when the previous stack exceeds some threshold.

Implement a data structure SetOfStacks that mimics this. SetOfStacks should be composed of several stacks, and should create a new stack once the previous one exceeds capacity. SetOfStacks.push() and SetOfStacks.pop() should behave identically to a single stack (that is, pop() should return the same values as it would if there were just a single stack).

Follow up: Implement a function popAt(int index) which performs a pop operation on a specific sub-stack.

Examples

Example 1

1
2
3
4


Input: capacity = 3
Operations: push(1), push(2), push(3), push(4), push(5)
Output: Two stacks - Stack 1: [1,2,3], Stack 2: [4,5]
Explanation: After pushing 3 elements, capacity is reached, so a new stack is created.

Example 2

1
2
3
4
5
6


Input: capacity = 2
Operations: push(1), push(2), push(3), pop(), popAt(0)
Output: Stack state after operations
Initial: Stack 1: [1,2], Stack 2: [3]
After pop(): Stack 1: [1,2], Stack 2: []
After popAt(0): Stack 1: [1], Stack 2: []

Solutions

Method 1 – Array-Based Implementation

Intuition

We need to maintain multiple stacks where each has a fixed capacity. When the current stack reaches capacity, we create a new one. The key operations are:

Push: Add to the last stack if it has space, otherwise create a new stack
Pop: Remove from the last stack, and remove empty stacks
PopAt: Remove from a specific stack by index

Approach

Use ArrayList to store multiple stacks: Each stack is implemented using a simple array
Track capacity per stack: Each individual stack has a maximum capacity
Handle push operations: Check if last stack has space, create new if needed
Handle pop operations: Remove from last stack, clean up empty stacks
Handle popAt operations: Allow popping from specific stack indices

Code

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108


#include <vector>
#include <stdexcept>

class Stack {
private:
    std::vector<int> data;
    int capacity;
    
public:
    Stack(int cap) : capacity(cap) {}
    
    void push(int value) {
        if (isFull()) {
            throw std::runtime_error("Stack is full");
        }
        data.push_back(value);
    }
    
    int pop() {
        if (isEmpty()) {
            throw std::runtime_error("Stack is empty");
        }
        int value = data.back();
        data.pop_back();
        return value;
    }
    
    bool isEmpty() const {
        return data.empty();
    }
    
    bool isFull() const {
        return data.size() >= capacity;
    }
    
    int size() const {
        return data.size();
    }
    
    int top() const {
        if (isEmpty()) {
            throw std::runtime_error("Stack is empty");
        }
        return data.back();
    }
};

class SetOfStacks {
private:
    std::vector<Stack> stacks;
    int capacity;
    
public:
    SetOfStacks(int cap) : capacity(cap) {}
    
    void push(int value) {
        if (stacks.empty() || stacks.back().isFull()) {
            stacks.emplace_back(capacity);
        }
        stacks.back().push(value);
    }
    
    int pop() {
        if (stacks.empty()) {
            throw std::runtime_error("No elements to pop");
        }
        
        int value = stacks.back().pop();
        
        // Remove empty stacks
        if (stacks.back().isEmpty()) {
            stacks.pop_back();
        }
        
        return value;
    }
    
    int popAt(int index) {
        if (index < 0 || index >= stacks.size()) {
            throw std::runtime_error("Invalid stack index");
        }
        
        int value = stacks[index].pop();
        
        // Remove empty stack if it's the last one
        if (stacks[index].isEmpty() && index == stacks.size() - 1) {
            stacks.pop_back();
        }
        
        return value;
    }
    
    bool isEmpty() const {
        return stacks.empty();
    }
    
    int size() const {
        int total = 0;
        for (const auto& stack : stacks) {
            total += stack.size();
        }
        return total;
    }
    
    int getStackCount() const {
        return stacks.size();
    }
};

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121


import java.util.*;

class Stack {
    private List<Integer> data;
    private int capacity;
    
    public Stack(int capacity) {
        this.capacity = capacity;
        this.data = new ArrayList<>();
    }
    
    public void push(int value) {
        if (isFull()) {
            throw new IllegalStateException("Stack is full");
        }
        data.add(value);
    }
    
    public int pop() {
        if (isEmpty()) {
            throw new IllegalStateException("Stack is empty");
        }
        return data.remove(data.size() - 1);
    }
    
    public boolean isEmpty() {
        return data.isEmpty();
    }
    
    public boolean isFull() {
        return data.size() >= capacity;
    }
    
    public int size() {
        return data.size();
    }
    
    public int top() {
        if (isEmpty()) {
            throw new IllegalStateException("Stack is empty");
        }
        return data.get(data.size() - 1);
    }
}

class SetOfStacks {
    private List<Stack> stacks;
    private int capacity;
    
    public SetOfStacks(int capacity) {
        this.capacity = capacity;
        this.stacks = new ArrayList<>();
    }
    
    public void push(int value) {
        if (stacks.isEmpty() || stacks.get(stacks.size() - 1).isFull()) {
            stacks.add(new Stack(capacity));
        }
        stacks.get(stacks.size() - 1).push(value);
    }
    
    public int pop() {
        if (stacks.isEmpty()) {
            throw new IllegalStateException("No elements to pop");
        }
        
        Stack lastStack = stacks.get(stacks.size() - 1);
        int value = lastStack.pop();
        
        // Remove empty stacks
        if (lastStack.isEmpty()) {
            stacks.remove(stacks.size() - 1);
        }
        
        return value;
    }
    
    public int popAt(int index) {
        if (index < 0 || index >= stacks.size()) {
            throw new IllegalArgumentException("Invalid stack index");
        }
        
        Stack stack = stacks.get(index);
        int value = stack.pop();
        
        // Remove empty stack if it's the last one
        if (stack.isEmpty() && index == stacks.size() - 1) {
            stacks.remove(index);
        }
        
        return value;
    }
    
    public boolean isEmpty() {
        return stacks.isEmpty();
    }
    
    public int size() {
        return stacks.stream().mapToInt(Stack::size).sum();
    }
    
    public int getStackCount() {
        return stacks.size();
    }
    
    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < stacks.size(); i++) {
            sb.append("Stack ").append(i).append(": ");
            Stack stack = stacks.get(i);
            sb.append("[");
            for (int j = 0; j < stack.size(); j++) {
                if (j > 0) sb.append(", ");
                // Note: This requires access to internal data, simplified for demo
            }
            sb.append("]\n");
        }
        return sb.toString();
    }
}

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103


class Stack:
    def __init__(self, capacity: int):
        self.capacity = capacity
        self.data = []
    
    def push(self, value: int) -> None:
        if self.is_full():
            raise Exception("Stack is full")
        self.data.append(value)
    
    def pop(self) -> int:
        if self.is_empty():
            raise Exception("Stack is empty")
        return self.data.pop()
    
    def is_empty(self) -> bool:
        return len(self.data) == 0
    
    def is_full(self) -> bool:
        return len(self.data) >= self.capacity
    
    def size(self) -> int:
        return len(self.data)
    
    def top(self) -> int:
        if self.is_empty():
            raise Exception("Stack is empty")
        return self.data[-1]

class SetOfStacks:
    def __init__(self, capacity: int):
        self.capacity = capacity
        self.stacks = []
    
    def push(self, value: int) -> None:
        # Create new stack if no stacks exist or last stack is full
        if not self.stacks or self.stacks[-1].is_full():
            self.stacks.append(Stack(self.capacity))
        
        self.stacks[-1].push(value)
    
    def pop(self) -> int:
        if not self.stacks:
            raise Exception("No elements to pop")
        
        value = self.stacks[-1].pop()
        
        # Remove empty stacks
        if self.stacks[-1].is_empty():
            self.stacks.pop()
        
        return value
    
    def pop_at(self, index: int) -> int:
        if index < 0 or index >= len(self.stacks):
            raise Exception("Invalid stack index")
        
        value = self.stacks[index].pop()
        
        # Remove empty stack if it's the last one
        if self.stacks[index].is_empty() and index == len(self.stacks) - 1:
            self.stacks.pop()
        
        return value
    
    def is_empty(self) -> bool:
        return len(self.stacks) == 0
    
    def size(self) -> int:
        return sum(stack.size() for stack in self.stacks)
    
    def get_stack_count(self) -> int:
        return len(self.stacks)
    
    def __str__(self) -> str:
        result = []
        for i, stack in enumerate(self.stacks):
            result.append(f"Stack {i}: {stack.data}")
        return "\n".join(result)

# Example usage and testing
def test_set_of_stacks():
    # Test with capacity 3
    stacks = SetOfStacks(3)
    
    # Push elements
    for i in range(1, 8):
        stacks.push(i)
        print(f"Pushed {i}, Stack count: {stacks.get_stack_count()}, Total size: {stacks.size()}")
    
    print(f"\nCurrent state:\n{stacks}")
    
    # Test pop
    print(f"\nPopped: {stacks.pop()}")
    print(f"Popped: {stacks.pop()}")
    
    # Test popAt
    print(f"PopAt(0): {stacks.pop_at(0)}")
    
    print(f"\nFinal state:\n{stacks}")

if __name__ == "__main__":
    test_set_of_stacks()

Complexity

⏰ Time complexity: O(1) for push, pop, and popAt operations
🧺 Space complexity: O(n), where n is the total number of elements stored across all stacks

Method 2 – Linked List Based Implementation with Rollover

Advanced Intuition

For the follow-up challenge, we might want to implement a more sophisticated version where popAt operations cause elements to “roll over” from subsequent stacks to maintain consistent stack sizes. This approach ensures that we don’t have gaps in the middle stacks.

Enhanced Approach

Implement rollover mechanism: When we pop from a middle stack, shift elements from subsequent stacks
Maintain consistent structure: Ensure all non-last stacks are at capacity
Use linked structure: Each stack node points to the next for easier rollover

Rollover Implementation

Python

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126


class Node:
    def __init__(self, value: int):
        self.value = value
        self.above = None
        self.below = None

class StackNode:
    def __init__(self, capacity: int):
        self.capacity = capacity
        self.size = 0
        self.top = None
        self.bottom = None
    
    def is_full(self) -> bool:
        return self.size == self.capacity
    
    def is_empty(self) -> bool:
        return self.size == 0
    
    def push(self, value: int) -> bool:
        if self.is_full():
            return False
        
        node = Node(value)
        if self.size == 0:
            self.bottom = node
        else:
            self.top.above = node
            node.below = self.top
        
        self.top = node
        self.size += 1
        return True
    
    def pop(self) -> int:
        if self.is_empty():
            raise Exception("Stack is empty")
        
        value = self.top.value
        self.top = self.top.below
        if self.top:
            self.top.above = None
        else:
            self.bottom = None
        
        self.size -= 1
        return value
    
    def remove_bottom(self) -> int:
        if self.is_empty():
            raise Exception("Stack is empty")
        
        value = self.bottom.value
        self.bottom = self.bottom.above
        if self.bottom:
            self.bottom.below = None
        else:
            self.top = None
        
        self.size -= 1
        return value

class SetOfStacksWithRollover:
    def __init__(self, capacity: int):
        self.capacity = capacity
        self.stacks = []
    
    def push(self, value: int) -> None:
        if not self.stacks or self.stacks[-1].is_full():
            self.stacks.append(StackNode(self.capacity))
        
        self.stacks[-1].push(value)
    
    def pop(self) -> int:
        if not self.stacks:
            raise Exception("No elements to pop")
        
        value = self.stacks[-1].pop()
        
        if self.stacks[-1].is_empty():
            self.stacks.pop()
        
        return value
    
    def pop_at(self, index: int) -> int:
        if index < 0 or index >= len(self.stacks):
            raise Exception("Invalid stack index")
        
        value = self.stacks[index].pop()
        self._shift_left(index)
        return value
    
    def _shift_left(self, index: int) -> None:
        """Shift elements left to fill gaps after popAt"""
        for i in range(index, len(self.stacks) - 1):
            if not self.stacks[i + 1].is_empty():
                # Move bottom element from next stack to current stack
                value = self.stacks[i + 1].remove_bottom()
                self.stacks[i].push(value)
        
        # Remove last stack if it becomes empty
        if self.stacks and self.stacks[-1].is_empty():
            self.stacks.pop()
    
    def is_empty(self) -> bool:
        return len(self.stacks) == 0
    
    def size(self) -> int:
        return sum(stack.size for stack in self.stacks)

# Example usage
def test_rollover_stacks():
    stacks = SetOfStacksWithRollover(3)
    
    # Push elements 1-9
    for i in range(1, 10):
        stacks.push(i)
    
    print(f"Initial size: {stacks.size()}")
    
    # Pop from middle stack - should cause rollover
    print(f"PopAt(1): {stacks.pop_at(1)}")
    print(f"Size after popAt: {stacks.size()}")

if __name__ == "__main__":
    test_rollover_stacks()

Rollover Complexity

⏰ Time complexity: O(1) for push and pop, O(k) for popAt where k is the number of stacks after the target index
🧺 Space complexity: O(n), where n is the total number of elements

Analysis

Trade-offs between approaches:

Method 1 (Simple): Easier to implement, popAt doesn’t maintain stack consistency
Method 2 (Rollover): More complex but maintains consistent stack structure after popAt operations

When to use each approach:

Use Method 1 when you need simple stack-of-stacks behavior and don’t care about gaps after popAt
Use Method 2 when you need to maintain consistent stack structure and prevent fragmentation

The choice depends on whether maintaining consistent stack sizes is important for your use case.

Problem#

Examples#

Example 1#

Example 2#

Solutions#

Method 1 – Array-Based Implementation#

Intuition#

Approach#

Code#

Complexity#

Method 2 – Linked List Based Implementation with Rollover#

Advanced Intuition#

Enhanced Approach#

Rollover Implementation#

Python#

Rollover Complexity#

Analysis#