Grokking#198

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:

    // Level cua tree dc danh so tu 0
    int depth(TreeNode* root) {
        int d = 0;
        while (root->left != nullptr) {
            d += 1;
            root = root->left;
        }
        return d;
    }

    // Return true if node idx at last level exist
    // O(d) time complexity
    bool exist(int idx, int d, TreeNode* root) {
        int left = 0, right = pow(2,d) -1;
        int pivot;
        for (int i = 0; i < d; ++i) {
            pivot = left + (right -left) /2;
            if( idx <= pivot) {
                root = root->left;
                right = pivot;
            } else {
                root = root->right;
                left = pivot;
            }
        }
        return root != nullptr;
    }

    int countNodes(TreeNode* root) {
        if (root == nullptr) return 0;

        int d = depth(root);
        if (d == 0) return 1;

        int left = 1, right = pow(2, d) -1; // Node 0 luon  ton tai nen ta danh so left bat dau tu 1
        int pivot;
        while (left <= right) {
            pivot = left + (right - left)/2;
            if (exist(pivot, d, root)) left = pivot + 1;
            else right = pivot -1;
        }

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