3.3 【発展問題】任意の部分スタックからpopする関数popAt(int index)を実装してください。

SetOfStacksTest.java

import java.util.*;
import org.junit.*;
import static org.junit.Assert.*;

public class SetOfStacksTest {
  @Test
  public void testBasicPushAndPop(){
    SetOfStacks set = new SetOfStacks();
    set.push(1);
    set.push(2);
    set.push(3);
    set.push(4);
    try {
      assertEquals(Integer.valueOf(4), set.pop());
      assertEquals(Integer.valueOf(3), set.pop());
      assertEquals(Integer.valueOf(2), set.pop());
      assertEquals(Integer.valueOf(1), set.pop());
    } catch (Exception e) {
      fail(e.getMessage());
    }
  }

  @Test
  public void testSetOfStacks(){
    SetOfStacks set = new SetOfStacks();
    assertEquals(1, set.stacks.size());
    set.push(1);
    set.push(2);
    set.push(3);
    assertEquals(1, set.stacks.size());
    set.push(4);
    assertEquals(2, set.stacks.size());
    set.push(5);
    set.push(6);
    set.push(7);
    assertEquals(3, set.stacks.size());
    try {
      assertEquals(Integer.valueOf(7), set.pop());
    } catch (Exception e){
      fail(e.getMessage());
    }
    assertEquals(2, set.stacks.size());
  }

  @Test
  public void testSetOfStacks3(){
    SetOfStacks set = new SetOfStacks();
    set.push(1);
    set.push(2);
    set.push(3);
    set.push(4);
    set.push(5);
    set.push(6);
    try {
      assertEquals(Integer.valueOf(3) ,set.popAt(0));
      assertEquals(Integer.valueOf(4) ,set.popAt(0));
      assertEquals(Integer.valueOf(6) ,set.pop());
    } catch (Exception e){
      fail(e.getMessage());
    }
  }

  @Test
  public void testSetOfStacks4(){
    SetOfStacks set = new SetOfStacks();
    set.push(1);
    set.push(2);
    set.push(3);
    assertEquals(0 ,set.stackIndex);
    set.push(4);
    set.push(5);
    set.push(6);
    assertEquals(1 ,set.stackIndex);
  }
}

class SetOfStacks {
  public List<Stack> stacks;
  public Stack currentStack;
  public int stackIndex;

  class Stack {
    /*
    This stack is supposed to have only 3 nodes.
    */
    public static final int CAPACITY = 3;
    public Node head;
    public Node tail;

    public Stack(){
      this.head = new Node();
      this.tail = new Node();
      this.head.next = this.tail;
    }

    public void push(Integer v){
      Node node = new Node();
      node.value = v;
      Node tmp = this.head.next;
      this.head.next = node;
      node.next = tmp;
    }

    public Integer pop() throws Exception {
      if(this.head.next == this.tail){
        throw new Exception("No node to be popped!");
      }
      Node node = this.head.next;
      Node tmp = node.next;
      this.head.next = tmp;
      return node.value;
    }

    public Integer getLast(){
      Node node = this.head;
      Node previous = node;
      while(node.next != this.tail){
        previous = node;
        node = node.next;
      }
      previous.next = this.tail;
      return node.value;
    }

    public boolean isFull(){
      if(this.length() >= CAPACITY){
        return true;
      }
      return false;
    }

    public boolean isEmpty(){
      if(this.head.next == this.tail){
        return true;
      }
      return false;
    }

    public int length(){
      int count = 0;
      Node node = this.head;
      while(node.next != this.tail){
        node = node.next;
        count++;
      }
      return count;
    }
  }

  class Node{
    public Node next = null;
    public Integer value = null;
  }

  public SetOfStacks(){
    this.stacks = new ArrayList<>();
    this.currentStack = new Stack();
    this.stacks.add(currentStack);
    stackIndex = 0;
  }

  public void push(Integer v){
    if(this.currentStack.isFull()){
      Stack stack = new Stack();
      stack.push(v);
      this.stacks.add(stack);
      this.currentStack = stack;
      stackIndex++;
    } else {
      this.currentStack.push(v);
    }
  }

  public Integer pop() throws Exception {
    Integer value;
    if(currentStack.isEmpty()){
      throw new Exception("No node to be popped.");
    }
    try {
      value = currentStack.pop();
    } catch (Exception e){
      value = null;
    }
    if(currentStack.isEmpty() && this.stacks.size() > 1){
      this.stacks.remove(this.stackIndex);
      this.stackIndex--;
      this.currentStack = this.stacks.get(stackIndex);
    }
    return value;
  }

  public Integer popAt(int i) throws Exception {
    if(i >= this.stacks.size()){
      throw new Exception("Index is out of range.");
    }

    Stack stack = this.stacks.get(i);
    Integer value = null;
    try {
      value = stack.pop();
    } catch (Exception e){
      throw e;
    }
    this.shift(i);
    return value;
  }

  public void shift(int index){
    if(index == stackIndex){
      return;
    }

    for(int i = index; i < this.stackIndex; i++){
      Stack stack = this.stacks.get(i);
      Stack stackBehind = this.stacks.get(i+1);
      Integer value = null;
      try {
        value = stackBehind.getLast();
      } catch (Exception e){
        System.out.println(e);
      }
      stack.push(value);
      if(stackBehind.isEmpty()){
        this.stacks.remove(i+1);
        stackIndex--;
        this.currentStack = this.stacks.get(stackIndex);
      }
    }
  }
}