Map implement by treap

Test.java

import java.util.Random;
import java.util.TreeMap;

/**
 * @author leer
 * Created at 12/17/18 6:31 PM
 */
public class Test {
  @org.junit.Test
  public void testTreap() {
    System.out.println("----------TreapMap test start----------");
    TreapMap<Integer, String> treap = new TreapMap<>();
    treap.put(1, "hello");
    treap.put(2, "world");
    System.out.println("treap[1] = " + treap.get(1));

    long time = 0;
    Random random = new Random();
    for (int i = 0; i < 1e6; i++) {
      long start = System.currentTimeMillis();
      treap.put(i, i + "test" + random.nextInt(Integer.MAX_VALUE));
      if (i % 5 == 0) {
        treap.remove(i);
      }
      time += (System.currentTimeMillis() - start);
    }

    System.out.println("size: " + treap.getSize());
    System.out.println("treap[1] = " + treap.get(4));
    System.out.println("TreapMap time: " + time);
    System.out.println("----------TreapMap test end----------");
  }

  @org.junit.Test
  public void testDelete() {
    TreapMap<Integer, String> treap = new TreapMap<>();
    treap.put(1, "hello");
    treap.put(2, "world");
    System.out.println("treap[1] = " + treap.get(1));

    Random random = new Random();
    for (int i = 0; i < 100; i++) {
      treap.put(i, i + "test" + random.nextInt(Integer.MAX_VALUE));
    }

    System.out.println("size: " + treap.getSize());
    treap.remove(1);
    System.out.println("After remove: treap[1] = " + treap.get(1));
    // test remove not exists key
    treap.remove(101);
    System.out.println("After remove size: " + treap.getSize());
    org.junit.Assert.assertEquals(treap.getSize(), 99);
  }

  @org.junit.Test
  public void testTreeMap() {
    System.out.println("----------TreeMap test start----------");
    TreeMap<Integer, String> treeMap = new TreeMap<>();
    treeMap.put(1, "hello");
    treeMap.put(2, "world");
    System.out.println("treeMap[1] = " + treeMap.get(1));

    Random random = new Random();
    long time = 0;
    for (int i = 0; i < 1e6; i++) {
      long start = System.currentTimeMillis();
      treeMap.put(i, i + "test" + random.nextInt(Integer.MAX_VALUE));
      if (i % 5 == 0) {
        treeMap.remove(i);
      }
      time += (System.currentTimeMillis() - start);
    }

    System.out.println("size: " + treeMap.size());
    System.out.println("treeMap[1] = " + treeMap.get(1));
    System.out.println("TreeMap time: " + time);
    System.out.println("----------TreeMap test end----------");
  }
}

TreapMap.java

import java.util.Iterator;
import java.util.Random;

/**
 * @author leer
 * Created at 12/17/18 4:35 PM
 * Tree and Heap
 */
@SuppressWarnings("unchecked")
public class TreapMap<K extends Comparable, V> implements Iterable {

  private Node<K, V> root;
  private int size;
  private int modeTime;

  private Random random;

  public TreapMap() {
    random = new Random(System.currentTimeMillis());
    root = null;
  }

  @Override
  public Iterator iterator() {
    throw new NotImplementedException();
  }

  static class Node<K extends Comparable, V> implements Comparable<Node> {
    Node[] children = new Node[2]; // left child, right child
    int priority;
    K key;
    V value;

    public Node(int priority, K k, V v) {
      this.priority = priority;
      this.key = k;
      this.value = v;
      children[0] = children[1] = null;
    }

    Node getLeftChild() {
      return children[0];
    }

    Node getRightChild() {
      return children[1];
    }

    int directionFor(K anotherKey) {
      if (anotherKey == this.key || anotherKey.equals(this.key)) return -1;
      // compareTo return negative integer when this object is less than another,
      // that is another is greater then this, so it should be in right child(children[1])
      return this.key.compareTo(anotherKey) < 0 ? 1 : 0;
    }

    @Override
    public int compareTo(Node another) {
      return this.key.compareTo(another.key) < 0 ? 0 : 1;
    }
  }

  private int getRandomPriority() {
    return random.nextInt(Integer.MAX_VALUE);
  }

  /**
   * d ^ 1 equals 1 -d, when d must be 1 or 0.
   * So this method handles both left rotate and right rotate
   *
   * @param parent parent node to rotate
   * @param d      rotate direction
   * @return parent node after rotate
   */
  private Node rotate(Node parent, int d) {
    if (d != 1 && d != 0) throw new ArrayIndexOutOfBoundsException();

    Node k = parent.children[d ^ 1];
    parent.children[d ^ 1] = k.children[d];
    k.children[d] = parent;
    parent = k;
    return parent;
  }

  private Node<K, V> insert(Node parent, K key, V value) {
    if (parent == null) {
      parent = new Node<>(getRandomPriority(), key, value);
      modeTime++;
    } else {
      int d = parent.directionFor(key);
      // key重复，直接覆盖value
      if (d == -1) {
        parent.value = value;
      } else {
        parent.children[d] = insert(parent.children[d], key, value);

        if (parent.children[d].priority > parent.priority) {
          // 插入后，如果子节点的优先级大于父节点的优先级，进行旋转
          parent = rotate(parent, d ^ 1);
          modeTime++;
        }
      }
    }
    return parent;
  }

  public void put(K key, V value) {
    int oldMod = modeTime;
    root = insert(root, key, value);
    if (oldMod != modeTime) size++;
  }

  private Node<K, V> delete(Node parent, K key) {
    V v = get(key);
    // key not exists
    if (v == null) {
      return parent;
    }

    int d = parent.directionFor(key);
    if (d == -1) {
      // 待删除的节点只有一个子节点或没有子节点
      if (parent.getLeftChild() == null) {
        parent = parent.getRightChild();
        modeTime++;
      } else if (parent.getRightChild() == null) {
        parent = parent.getLeftChild();
        modeTime++;
      } else {
        //待删除的节点有两个子节点，把优先级高的子节点旋转的父节点，再递归删除待删除的节点
        //如果左节点的优先级高于右节点，右旋，接着在右子树中递归删除。
        int d2 = parent.getLeftChild().compareTo(parent.getRightChild());
        parent = rotate(parent, d2);
        delete(parent.children[d2], key);
      }
    } else {
      parent = delete(parent.children[d], key);
    }
    return parent;
  }

  public void remove(K key) {
    int oldMod = modeTime;
    root = delete(root, key);
    if (modeTime != oldMod) size--;
  }

  public V get(K key) {
    Node<K, V> p = root;
    while (p != null) {
      int d = p.directionFor(key);
      if (d == -1) return p.value;
      else p = p.children[d];
    }
    return null;
  }

  public int getSize() {
    return size;
  }
}