#ctci

SetOfStacks.java

import java.util.ArrayList;
import java.util.Stack;

public class SetOfStacks<E>{
    private int numOfStacks;
    private ArrayList<Stack<E>> stacks;
    private int capacity;
    
    public SetOfStacks(int capacity){
        this.capacity = capacity;
        stacks = new ArrayList<Stack<E>>();
        Stack<E> s = new Stack<E>();
        numOfStacks++;
    }
    
    public void push(E val){
        Stack<E> s = stacks.get(numOfStacks-1);
        if(s.size() == capacity){
            Stack<E> newS = new Stack<E>();
            s = newS;
            stacks.add(newS);
            numOfStacks++;
        }
        s.push(val);
    }
    
    public E pop() throws Exception{
        if(numOfStacks == 0) throw new Exception("Stack underflow!");
        
        Stack<E> s = stacks.get(numOfStacks-1);
        E val = s.pop();
        if(s.size() == 0){
            stacks.remove(numOfStacks-1);
            numOfStacks--;
        }
        return E;
    }
    
    public E popAt(int index) throws Exception{
        if(numOfStacks == 0) throw new Exception("Stack underflow!");
        if(numOfStacks <= index) throw new Exception("No such stack exists!");
        
        Stack<E> s = stacks.get(index);
        E val = s.pop();
        if(s.size() == 0){
            stacks.remove(index);
            numOfStacks--;
        }
        return E;
    }
    
}