CTCI 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 structureSetOf Stacks that mimics this. SetOf Stacks should be composed of several stacks and should create a new stack once the previous one exceeds capacity. SetOfStacks.push() and SetOfStacks.pop() should be have 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.

Using arraylist could easily implement stacks, and use another arraylist to manage these stacks, when the current stack is full, add a new one, when the current is empty, remove it from the list. In terms of popAt(), I assume that after pop the top from the specific stack, the hole would not be filled in the following push process. We could ask interviewer to clarify this problem. Besides, I assume the data type to be integer, we could use generalization to implement the class.

/*Use arraylist to implement stacks and another arraylist to store these stacks, assume the elements
are integer to ease the discussion. I assume that after a popAt() call, the newly pushed elements should
still see the entire stack as a single stack.*/

import java.util.*;

public class SetofStacks {
    ArrayList<ArrayList<Integer>> sets;
    int stackSize; int stackNum = -1;
    public SetofStacks(int size) {
        sets = new ArrayList<ArrayList<Integer>>();
        stackSize = size;
    } 

    public void push(int key) {
        //If there is no stack yet or the current stack is full, create a new one
        if(stackNum == -1 || sets.get(stackNum).size()+1 > stackSize) {
            sets.add(new ArrayList<Integer>());
            stackNum++;
        }
        sets.get(stackNum).add(key);
    }

    public int pop() throws Exception {
        if(stackNum == -1) {
            throw new Exception("No Elements");
        }

        int top = sets.get(stackNum).size();
        int key = sets.get(stackNum).get(top-1);
        sets.get(stackNum).remove(top-1);

        //If a stack is empty, remove it from the list
        if(sets.get(stackNum).size() == 0) {
            sets.remove(stackNum);
            stackNum--;
        }

        return key;
    }

    public int popAt(int index) throws Exception{
        if(stackNum < index) {
            throw new Exception("No Such Stack");
        }

        int top = sets.get(index).size();
        int key = sets.get(index).get(top-1);
        sets.get(index).remove(top-1);

        return key;
    }

    public static void main(String[] args) {
        SetofStacks ss = new SetofStacks(2);
        try {
            ss.push(1); ss.push(2); ss.push(3);
            System.out.println(ss.pop() + " " + ss.pop() + " " +ss.pop());
            ss.push(1); ss.push(2); ss.push(3); ss.push(4);
            System.out.println(ss.popAt(0) + " " + ss.popAt(1));
            ss.push(5); ss.push(6);
            System.out.println(ss.pop() + " " + ss.pop() + " " +ss.pop());
        }
        catch(Exception e) {
            System.out.println(e.toString());
        }
    }
}

CTCI 3.3