Avl Tree with parent

IterableAvlTree.java

  private static class Entry<TKey extends Comparable<TKey>> {
    private Entry<TKey> avlParent = null;
    private Entry<TKey> avlRight = null;
    private Entry<TKey> avlLeft = null;
    private int avlHeight = 1;

    private final TKey key;
    public Entry(TKey key) {
      this.key = key;
    }

    public int compareKey(TKey cmpKey) {
      return key.compareTo(cmpKey);
    }

    public int compareTo(Entry<TKey> other) {
      return compareKey(other.key);
    }
  }

  private static class IterableAvlTree {
    public static <T extends Comparable<T>> Entry<T> get(Entry<T> root, T key) {
      while (root != null) {
        int cmp = root.compareKey(key);
        if (cmp > 0) {
          root = root.avlLeft;
        } else if (cmp < 0) {
          root = root.avlRight;
        } else {
          return root;
        }
      }
      return null;
    }

    public static <T extends Comparable<T>> Entry<T> getFirstEntry(Entry<T> root) {
      if (root != null) {
        while (root.avlLeft != null) {
          root = root.avlLeft;
        }
      }
      return root;
    }

    public static <T extends Comparable<T>> Entry<T> getLastEntry(Entry<T> root) {
      if (root != null) {
        while (root.avlRight != null) {
          root = root.avlRight;
        }
      }
      return root;
    }

    public static <T extends Comparable<T>> Entry<T> successor(Entry<T> p) {
      if (p == null) return null;
      if (p.avlRight != null) {
        Entry<T> q = p.avlRight;
        while (q.avlLeft != null) {
          q = q.avlLeft;
        }
        return q;
      } else {
        Entry<T> q = p.avlParent;
        while (q != null && p == q.avlRight) {
          p = q;
          q = q.avlParent;
        }
        return q;
      }
    }

    public static <T extends Comparable<T>> Entry<T> predecessor(Entry<T> p) {
      if (p == null) return null;
      if (p.avlLeft != null) {
        Entry<T> q = p.avlLeft;
        while (q.avlRight != null) {
          q = q.avlRight;
        }
        return q;
      } else {
        Entry<T> q = p.avlParent;
        while (q != null && p == q.avlLeft) {
          p = q;
          q = q.avlParent;
        }
        return q;
      }
    }

    public static <T extends Comparable<T>> Entry<T> insert(Entry<T> root, Entry<T> node) {
      if (root == null) return node;

      int cmp = root.compareTo(node);
      assert cmp != 0 : "Unsupported operation replace";
      if (cmp > 0) {
        root.avlLeft = insert(root.avlLeft, node);
      } else {
        root.avlRight = insert(root.avlRight, node);
      }
      return balance(root);
    }

    public static <T extends Comparable<T>> Entry<T> remove(Entry<T> root, T key) {
      if (root == null) return null;

      int cmp = root.compareKey(key);
      if (cmp == 0) {
        if (root.avlLeft == null && root.avlRight == null) {
          return null;
        } else if (root.avlLeft == null) {
          if (root.avlRight != null) {
            root.avlRight.avlParent = root.avlParent;
          }
          return root.avlRight;
        } else if (root.avlRight == null) {
          if (root.avlLeft != null) {
            root.avlLeft.avlParent = root.avlParent;
          }
          return root.avlLeft;
        } else {
          Entry<T> left = root.avlRight;
          while (left.avlLeft != null) {
            left = left.avlLeft;
          }
          left.avlRight = removeLeft(root.avlRight);
          left.avlLeft = root.avlLeft;
          return balance(left);
        }
      } else if (cmp > 0) {
        root.avlLeft = remove(root.avlLeft, key);
      } else {
        root.avlRight = remove(root.avlRight, key);
      }
      return balance(root);
    }

    private static <T extends Comparable<T>> Entry<T> removeLeft(Entry<T> node) {
      if (node.avlLeft == null) {
        node.avlParent = null;
        return node.avlRight;
      } else {
        node.avlLeft = removeLeft(node.avlLeft);
        return balance(node);
      }
    }

    private static <T extends Comparable<T>> Entry<T> rotateRight(Entry<T> node) {
      Entry<T> temp = node.avlLeft;
      node.avlLeft = temp.avlRight;
      if (node.avlLeft != null) {
        node.avlLeft.avlParent = node;
      }

      temp.avlRight = node;
      temp.avlParent = node.avlParent;
      node.avlParent = temp;
      if (temp.avlParent != null) {
        if (temp.avlParent.avlRight == node) {
          temp.avlParent.avlRight = temp;
        } else {
          temp.avlParent.avlLeft = temp;
        }
      }
      computeHeight(node);
      computeHeight(temp);
      return temp;
    }

    private static <T extends Comparable<T>> Entry<T> rotateLeft(Entry<T> node) {
      Entry<T> temp = node.avlRight;
      node.avlRight = temp.avlLeft;
      if (node.avlRight != null) {
        node.avlRight.avlParent = node;
      }

      temp.avlLeft = node;
      temp.avlParent = node.avlParent;
      node.avlParent = temp;
      if (temp.avlParent != null) {
        if (temp.avlParent.avlRight == node) {
          temp.avlParent.avlRight = temp;
        } else {
          temp.avlParent.avlLeft = temp;
        }
      }
      computeHeight(node);
      computeHeight(temp);
      return temp;
    }

    private static <T extends Comparable<T>> void computeHeight(Entry<T> node) {
      node.avlHeight = 1;
      if (node.avlLeft != null) {
        node.avlLeft.avlParent = node;
        node.avlHeight += node.avlLeft.avlHeight;
      }
      if (node.avlRight != null) {
        node.avlRight.avlParent = node;
        node.avlHeight += node.avlRight.avlHeight;
      }
    }

    private static <T extends Comparable<T>> int getBalance(Entry<T> node) {
      int balance = 0;
      if (node.avlLeft != null) {
        balance += node.avlLeft.avlHeight;
      }
      if (node.avlRight != null) {
        balance -= node.avlRight.avlHeight;
      }
      return balance;
    }

    private static <T extends Comparable<T>> Entry<T> balance(Entry<T> node) {
      computeHeight(node);
      int balance = getBalance(node);
      if (balance > 1) {
        if (getBalance(node.avlLeft) < 0) {
          node.avlLeft = rotateLeft(node.avlLeft);
        }
        return rotateRight(node);
      } else if (balance < -1) {
        if (getBalance(node.avlRight) > 0) {
          node.avlRight = rotateRight(node.avlRight);
        }
        return rotateLeft(node);
      }
      return node;
    }
  }

  private static void testAvl() {
    java.util.TreeMap<Long, Object> treeMap = new java.util.TreeMap<Long, Object>();
    Entry<Long> root = null;

    long st = System.currentTimeMillis();

    long seed = System.currentTimeMillis();
    System.out.println("seed=" + seed);
    final java.util.Random rand = new java.util.Random(seed);
    final int NELEM = 1 + rand.nextInt(5000);
    System.out.println("NELEM " + NELEM);

    System.out.println("INSERT");
    for (int i = 0; i < NELEM; ++i) {
      long key = rand.nextInt(NELEM);
      if (IterableAvlTree.get(root, key) != null) {
        i--;
        continue;
      }
      root = IterableAvlTree.insert(root, new Entry<Long>(key));
      treeMap.put(key, null);
      for (Long keyX: treeMap.keySet()) {
        Entry<Long> node = IterableAvlTree.get(root, keyX);
        assert node != null;
        assert node.key.equals(keyX);
      }
      traverse(root);
    }

    System.out.println("REMOVE");
    for (int i = 0; i < NELEM; ++i) {
      long key = rand.nextInt(NELEM);
      Entry<Long> node = IterableAvlTree.get(root, key);
      if (!treeMap.containsKey(key)) {
        assert node == null;
        continue;
      }
      treeMap.remove(key);
      assert node.key == key;
      root = IterableAvlTree.remove(root, key);
      for (Long keyX: treeMap.keySet()) {
        node = IterableAvlTree.get(root, keyX);
        assert node != null;
        assert node.key.equals(keyX);
      }
      traverse(root);
    }

    for (Long keyX: treeMap.keySet()) {
      root = IterableAvlTree.remove(root, keyX);
      traverse(root);
    }

    assert root == null : "tree is not empty";

    long et = System.currentTimeMillis();
    System.out.println(et - st);

    for (int i = 1; i <= 20; ++i) {
      Entry<Long> entry = new Entry<Long>((long)i);
      root = IterableAvlTree.insert(root, entry);
      traverse(root);
    }

    for (int i = 1; i <= 20; ++i) {
      traverse(root);
      //dumpOrder(root);
      root = IterableAvlTree.remove(root, (long)i);
    }
    traverse(root);
    assert root == null;

    for (int i = 0; i < NELEM; ++i) {
      Entry<Long> entry = new Entry<Long>((long)i);
      root = IterableAvlTree.insert(root, entry);
      traverse(root);
    }

    for (int i = NELEM - 1; i >= 0; --i) {
      root = IterableAvlTree.remove(root, (long)i);
      traverse(root);
    }
  }