2-3 Tree

https://www.youtube.com/watch?v=C3SsdUqasD4 presents an example of 2-3 tree. To help myself get a better understanding of this data structure, I tried to implement it, and built a set using it.

The following snippet contains the example used in the video, and testing code to show it behaves like a set (HashSet in Java Standard Library).

import java.util.*;

class Main {
  static class Tree23 {

    static class Node {
      enum Child {c1, c2, c3}

      int left, right;
      Node c1, c2, c3;
      Node parent;
      int count;

      Node(int e, Node parent) {
        left = e;
        this.parent = parent;
        count = 1;
      }

      boolean isLeaf() {
        return c1 == null;
      }

      void set(Child id, Node c) {
        switch (id) {
          case c1:
            this.c1 = c;
            break;
          case c2:
            this.c2 = c;
            break;
          case c3:
            this.c3 = c;
            break;
        }
        c.parent = this;
      }
    }

    Node root;

    boolean has(int e) {
      return has(e, root);
    }

    void add(int e) {
      if (root == null) {
        root = new Node(e, null);
        return;
      }
      add(e, root);
    }

    void print() {
      System.out.println("PRINTING:");
      if (root == null) {
        System.out.println("Empty tree.");
        return;
      }

      printLevelByLevel(Collections.singletonList(root));
    }

    private boolean has(int e, Node node) {
      if (node == null) return false;
      if (node.left == e || (node.count == 2 && node.right == e)) return true;
      if (node.count == 1) {
        if (e < node.left) {
          return has(e, node.c1);
        } else {
          return has(e, node.c3);
        }
      }

      assert node.count == 2;
      if (e < node.left) {
        return has(e, node.c1);
      } else if (e > node.right) {
        return has(e, node.c3);
      } else {
        return has(e, node.c2);
      }
    }

    private void promote(Node node) {
      assert node.count == 1;
      assert node.c1.count == 1;
      assert node.c3.count == 1;
      assert node.c2 == null;

      if (node.parent == null) {
        if (root != node) root = node;
        return;
      }

      assert node.parent.count != 0;

      if (node.parent.count == 1) {
        if (node.parent.left < node.left) {
          node.parent.right = node.left;
          node.parent.set(Node.Child.c2, node.c1);
          node.parent.set(Node.Child.c3, node.c3);
        } else {
          node.parent.right = node.parent.left;
          node.parent.left = node.left;
          node.parent.set(Node.Child.c1, node.c1);
          node.parent.set(Node.Child.c2, node.c3);
        }
        node.c1 = node.c3 = null;
        node.parent.count = 2;
        return;
      }

      Node parent = node.parent;

      assert parent.count == 2;
      assert parent.c1 != null && parent.c2 != null && parent.c3 != null;

      if (node.left < parent.left) {
        Node sibling = new Node(parent.right, parent);

        sibling.set(Node.Child.c1, parent.c2);
        sibling.set(Node.Child.c3, parent.c3);

        parent.c1 = node;
        parent.c2 = null;
        parent.c3 = sibling;
        // parent.left stays the same
        parent.count = 1;
      } else if (node.left > parent.right) {
        Node sibling = new Node(parent.left, parent);

        sibling.set(Node.Child.c1, parent.c1);
        sibling.set(Node.Child.c3, parent.c2);

        parent.c1 = sibling;
        parent.c2 = null;
        parent.c3 = node;
        parent.left = parent.right;
        parent.count = 1;
      } else {
        Node left_sibling = new Node(parent.left, parent);
        Node right_sibling = new Node(parent.right, parent);

        left_sibling.set(Node.Child.c1, parent.c1);
        left_sibling.set(Node.Child.c3, node.c1);

        right_sibling.set(Node.Child.c1, node.c3);
        right_sibling.set(Node.Child.c3, parent.c3);

        parent.c1 = left_sibling;
        parent.c2 = null;
        parent.c3 = right_sibling;
        parent.left = node.left;
        parent.count = 1;
      }
      promote(parent);
    }

    private void squeeze(Node node, int e) {
      assert node.isLeaf();
      if (node.count == 1) {
        if (node.left < e) {
          node.right = e;
        } else {
          node.right = node.left;
          node.left = e;
        }
        node.count = 2;
      } else {
        if (e < node.left) {
          node.c1 = new Node(e, node);
          node.c3 = new Node(node.right, node);
          // node.left stays the same
        } else if (e > node.right) {
          node.c1 = new Node(node.left, node);
          node.c3 = new Node(e, node);
          node.left = node.right;
        } else {
          node.c1 = new Node(node.left, node);
          node.c3 = new Node(node.right, node);
          node.left = e;
        }

        assert node.c2 == null;
        node.count = 1;
        promote(node);
      }
    }

    private void add(int e, Node node) {
      assert node.count == 1 || node.count == 2;
      if (node.left == e || (node.count == 2 && node.right == e)) return;
      if (node.isLeaf()) {
        squeeze(node, e);
      } else {
        if (node.count == 1) {
          if (node.left < e) {
            add(e, node.c3);
          } else {
            add(e, node.c1);
          }
        } else {
          if (e < node.left) {
            add(e, node.c1);
          } else if (e > node.right) {
            add(e, node.c3);
          } else {
            add(e, node.c2);
          }
        }
      }
    }

    private void printLevelByLevel(List<Node> nodes) {
      for (Node node : nodes) {
        print(node);
      }
      System.out.println("");

      if (nodes.get(0).isLeaf()) return;

      List<Node> new_nodes = new ArrayList<>();

      for (Node node : nodes) {
        assert node.c1 != null;
        assert node.c1.parent == node;

        if (node.count == 1) assert node.c2 == null;

        new_nodes.add(node.c1);
        if (node.c2 != null) {
          new_nodes.add(node.c2);
          assert node.c2.parent == node;
        }
        if (node.c3 != null) {
          new_nodes.add(node.c3);
          assert node.c3.parent == node;
        }

      }
      printLevelByLevel(new_nodes);
    }

    private void print(Node node) {
      assert node.count > 0;
      if (node.count == 1) {
        System.out.format("( %d )", node.left);
      } else {
        System.out.format("( %d, %d )", node.left, node.right);
      }
    }
  }

  private static void test() {
    Random rand = new Random();
    Tree23 t = new Tree23();
    Set<Integer> set = new HashSet<>();
    int size = 10_000_000;
    for (int i = 0; i < size; ++i) {
      if (rand.nextBoolean()) {
        t.add(i);
        set.add(i);
      }
    }

    for (int i = 0; i < size; ++i) {
      assert t.has(i) == set.contains(i);
    }
    System.out.println("Done Testing.");
  }

  private static void printTest() {
    Tree23 t = new Tree23();
    int inputs[] = {5, 10, 20, 8, 7, 15, 30, 11};
    // int inputs[] = {4, 1, 9, 8, 5, 6, 7};
    for (int input : inputs) {
      t.add(input);
    }
    t.print();
  }

  public static void main(String[] args) {
    printTest();
    // test();
  }
}