Friday, June 5, 2015

[LeetCode] Count Complete Tree Nodes

Count Complete Tree Nodes

Given a complete binary tree, count the number of nodes.
Definition of a complete binary tree from Wikipedia:
In a complete binary tree every level, except possibly the last, is completely filled, and all nodes in the last level are as far left as possible. It can have between 1 and 2nodes inclusive at the last level h.
Solution: Binary Search.
A naive solution would be simply traversing the whole tree, which takes O(n) time, where n is the total number of nodes in the binary tree.
Optimization: Suppose the last level is h, and we have t nodes in level h, where 1 <= t <= 2^h. The observation is that all the t nodes are aligned from the leftmost. That is, if we label the nodes in the last level h, then nodes 1, ..., t exist, while nodes t + 1, ...., 2^h do not exist. This can be visualized as an non-increasing sorted array: [1, ..., 1, 0, ..., 0]. Binary search helps avoid traversing the whole binary tree.
Complexity: O(log^2 n).
Note that n = O(2^h). We need to perform O(log 2^h) = O(h) binary search operations, and each binary search operation takes O(h) time (since the height of the binary tree is O(h)). Hence the total running time is O(h^2) = O(log^2 n).
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    private int countLevels(TreeNode root) {
        int h = -1;
        TreeNode p = root;
        while (p != null) {
            h++;
            p = p.left;
        }
        return h;
    }
    
    private boolean existInLastLevel(
        TreeNode root,
        int path_val, 
        int h
    ) {
        TreeNode p = root;
        int power = 1 << (h - 1);
        for (int i = 0; i < h; i++) {
            if (path_val <= power) {
                if (p.left == null) {
                    return false;
                }
                p = p.left;
            } else {
                path_val -= power;
                if (p.right == null) {
                    return false;
                }
                p = p.right;
            }
            power /= 2;
        }
        return true;
    }
    
    private boolean isRightmost(
        TreeNode root,
        int path_val,
        int h
    ) {
        if (!existInLastLevel(root, path_val, h)) {
            return false;
        }
        if (path_val == (1 << h)) {
            return true;
        }
        if (!existInLastLevel(root, path_val + 1, h)) {
            return true;
        }
        return false;
    }
    
    private boolean toTheLeftOfRightmost(
        TreeNode root,
        int path_val,
        int h
    ) {
        if (
            existInLastLevel(root, path_val, h) &&
            !isRightmost(root, path_val, h)
        ) {
            return true;
        }
        return false;
    }
    
    public int countNodes(TreeNode root) {
        if (root == null) {
            return 0;
        }
        int h = countLevels(root);
        int low = 1;
        int high = 1 << h;
        int cntInLastLevel = -1;
        while (low <= high) {
            int mid = low + (high - low) / 2;
            if (isRightmost(root, mid, h)) {
                cntInLastLevel = mid;
                break;
            }
            if (toTheLeftOfRightmost(root, mid, h)) {
                low = mid + 1;                
            } else {
                high = mid - 1;
            }
        }
        int cntAboveLastLevel = (1 << h) - 1;
        return cntInLastLevel + cntAboveLastLevel;
    }
}

No comments:

Post a Comment