Question3_3.java

import java.util.ArrayList;
import java.util.EmptyStackException;
import java.util.LinkedList;

/**
 * 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 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.
 * 
 * Time  Complexity: O(1) for pop, push
 *                   O(m) for popAt (m is the number of stacks)
 * Space Complexity: O(n) (n is the number of total elements)
 */

public class Question3_3<T> {
    private ArrayList<LinkedList<T>> setOfStacks;
    private int capacity;
    
    public Question3_3 (int capcity) {
        if (capcity <= 0) {
            throw new IllegalArgumentException();
        }
        setOfStacks = new ArrayList<LinkedList<T>>();
        this.capacity = capcity;
    }
    
    public void push(T element) {
        if (setOfStacks.isEmpty() || setOfStacks.get(setOfStacks.size() - 1).size() == capacity) {
            setOfStacks.add(new LinkedList<T>());
        }
        LinkedList<T> stack = setOfStacks.get(setOfStacks.size() - 1);
        stack.addLast(element);
    }
    
    public T pop() {
        if (setOfStacks.isEmpty()) {
            throw new EmptyStackException();
        }
        LinkedList<T> stack = setOfStacks.get(setOfStacks.size() - 1);
        T element = stack.removeLast();
        if (stack.isEmpty()) {
            setOfStacks.remove(setOfStacks.size() - 1);
        }
        return element;
    }
    
    public T popAt(int index) {
        if (index < 0 || index >= setOfStacks.size()) {
            throw new IllegalArgumentException();
        }
        
        if (index == setOfStacks.size() - 1) {
            return pop();
        }
        
        LinkedList<T> stack = setOfStacks.get(index);
        T element = stack.removeLast();
        rollover(index, index + 1);
        return element;
    }
    
    private void rollover(int curr, int next) {
        if (next == setOfStacks.size()) {
            if (setOfStacks.get(curr).isEmpty()) {
                setOfStacks.remove(curr);
            }
            return;
        }
        
        LinkedList<T> stack = setOfStacks.get(curr);
        LinkedList<T> nextStack = setOfStacks.get(next);
        stack.addLast(nextStack.removeFirst());
        rollover(next, next + 1);
    }
    
    public void printStacks() {
        for (LinkedList<T> stack : setOfStacks) {
            for (T element : stack) {
                System.out.print(element + "\t");
            }
            System.out.println();
        }
    }
    
    
    public static void main(String[] args) {
        Question3_3<Integer> question3_3 = new Question3_3<Integer>(3);
        for (int i = 0; i < 12; ++i) {
            question3_3.push(i);
        }
        System.out.println("Set of stacks:");
        question3_3.printStacks();
        
        System.out.println();
        for (int i = 0; i < 5; ++i) {
            System.out.println("pop: " + question3_3.pop());
        }
        System.out.println();
        System.out.println("Set of stacks:");
        question3_3.printStacks();
        
        System.out.println();
        System.out.println("pop: " + question3_3.popAt(0));
        System.out.println();
        System.out.println("Set of stacks:");
        question3_3.printStacks();
    }
}