Gyh's Braindump

222. Count Complete Tree Nodes

tags
Tree, BinarySearch, Recursion
source
leetcode

Edge Cases

Solution 1 - Recursion

NOTE: this can also apply to non-complete tree

class Solution {
    public int countNodes(TreeNode root) {
        return root == null ? 0 : 1 + countNodes(root.left) + countNodes(root.right);
    }
}

Complexity

  • time: O(N)
  • space: O(lgN)

key is to find number of leafs in the final level

use binary search to find first position missing: O(d) when doing binary search, we need to test if the node exists in each compare, use binary search again to go down the tree: O(d)

class Solution {
    public int countNodes(TreeNode root) {
        if (root == null) return 0;

        // compute depth
        int d = computeDepth(root);

        // binary search
        int left = 0, right = (int) Math.pow(2, d) - 1;
        while (left <= right) {
            int pivot = (left + right) / 2;
            if (exists(root, d, pivot)) left = pivot + 1;
            else right = pivot - 1;
        }

        return left + (int) Math.pow(2, d) - 1;
    }

    private int computeDepth(TreeNode root) {
        int d = 0;
        while (root != null) {
            root = root.left;
            d++;
        }
        return d - 1;
    }

    private boolean exists(TreeNode root, int d, int idx) {
        int left = 0, right = (int) Math.pow(2, d) - 1;
        for (int i = 0; i < d; i++) {
            int mid = (left + right) / 2;
            if (idx <= mid) {
                root = root.left;
                right = mid;
            } else {
                root = root.right;
                left = mid + 1;
            }
        }
        return root != null;
    }
}

Complexity

  • time: O(d^2)
  • space: O(1)