Three in One

Use a single array to implement three stacks.

Solution

The most simple solution involves dividing the array into three equal parts - one for each stack. As elements are pushed onto each stack, they are appended to the end of the appropriate array segment. We can use an index to keep track of the top of each stack - and when elements are pushed from the stack, we will remove them from the array and decrement the index.

class MultiStack:
  def __init__(self, num_stacks=3, capacity=9):
    self.num_stacks = num_stacks
    self.capacity = capacity
    self.values = [None]*capacity
    self.sizes = [0]*num_stacks

  def push(self, stack_num, value):
    size = self.sizes[stack_num - 1]
    offset = self.capacity // 3 * (stack_num - 1)
    i = offset + size

    # add to array
    self.values[i] = value
    self.sizes[stack_num - 1] += 1

  def pop(self, stack_num):
    size = self.sizes[stack_num - 1]
    offset = self.capacity // 3 * (stack_num - 1)
    i = offset + size

    # remove from array
    value = self.values[i - 1]
    self.values[i - 1] = None

    self.sizes[stack_num - 1] -= 1

    return value