Complete Binary Tree Inserter LeetCode solution in Java

1
class CBTInserter {
2

3
  private TreeNode root;
4
  private int total;
5

6
  private int count(TreeNode root) {
7
    if (root == null) return 0;
8
    return 1 + count(root.left) + count(root.right);
9
  }
10

11
  public CBTInserter(TreeNode root) {
12
    this.root = root;
13
    total = count(root);
14
  }
15

16
  private int insertBinary(int val, int k, int right) {
17
    int left = 0;
18
    var ptr = root;
19
    while (left < right) {
20
      if (left == right - 1) {
21
        if (ptr.left == null) ptr.left = new TreeNode(val);
22
        else ptr.right = new TreeNode(val);
23
        return ptr.val;
24
      }
25
      int mid = (right - left) / 2 + left;
26
      if (mid >= k) {
27
        ptr = ptr.left;
28
        right = mid;
29
      } else if (mid < k) {
30
        left = mid + 1;
31
        ptr = ptr.right;
32
      }
33
    }
34
    return 0;
35
  }
36

37
  public int insert(int val) {
38
    int depth = 0;
39
    int n = total;
40
    while (n > 0) {
41
      depth++;
42
      n /= 2;
43
    }
44
    if ((1 << depth) - 1 == total) depth++;
45
    // k is the new index in the lowest level of tree, right is the max index of the lowest level of
46
    // tree
47
    // e.g. for tree
48
    // 1
49
    // / \
50
    // 2 new
51
    // index 0 1
52
    // it's insertBinary(val, 1, 1)
53
    var res = insertBinary(val, total - (1 << (depth - 1)) + 1, (1 << (depth - 1)) - 1);
54
    total++;
55
    return res;
56
  }
57

58
  public TreeNode get_root() {
59
    return root;
60
  }
61
}