Chapter 3 Stacks and Queues - 3.3

Problem 3.3: 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.

It is a quite easy one. However, the requirement of follow up question is not clear and we should discuss details with interviewers.

from stack import *
MAX_CAPACITY = 3
class SetOfStacks:
    def __init__(self):
        self.stacks = [stack()]
        self.last = 0
    def push(self, value):
        if self.stacks[self.last].size() >= MAX_CAPACITY:
            self.stacks.append(stack())
            self.last = self.last + 1
        self.stacks[self.last].push(value)
    def pop(self):
        if self.stacks[self.last].size() <= 0:
            if len(self.stacks) <= 1:
                return False
            del self.stacks[self.last]
            self.last = self.last - 1
        return self.stacks[self.last].pop()
    def popAt(self, index):
        if (index < 0) or (index > self.last):
            return False
        to_return = self.stacks[index].pop()
        if self.stacks[index].size() <= 0:
            del self.stacks[index]
            self.last = self.last - 1
        return to_return


# Test cases
if __name__ == "__main__":
    stacks_set = SetOfStacks()
    for i in range(0, 10):
        print "push", i
        stacks_set.push(i)
    for i in range(0, 3):
        print "pop", stacks_set.pop()
    for i in range(0, 5):
        print "push", i
        stacks_set.push(i)
    for i in range(0, 13):
        print "popAt(0)", stacks_set.popAt(0)