222 - Count Complete Tree Nodes

#medium

Given the root of a complete binary tree, return the number of the nodes in the tree.

According to Wikipedia, every level, except possibly the last, is completely filled in a complete binary tree, and all nodes in the last level are as far left as possible. It can have between 1 and 2h nodes inclusive at the last level h.

Design an algorithm that runs in less than O(n) time complexity.

Example 1:

Input: root = [1,2,3,4,5,6]
Output: 6

Example 2:

Input: root = []
Output: 0

Example 3:

Input: root = [1]
Output: 1

思路

子樹是一個complete binary tree = 左右高度相同。只走最外側節點，中間節點皆不訪問。

• 時間複雜度：O(log n × log n)
• 空間複雜度：O(log n)

Solution

/**
 * 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:
    int countNodes(TreeNode* root) {
        if ( root == NULL ) return 0;
        TreeNode * left = root;
        TreeNode * right = root;
        int leftDepth = 0, rightDepth = 0;
        // 判斷子樹是完全樹的狀況
        while ( left != NULL ) { // 求左子樹深度
            left = left -> left;
            leftDepth++;
        }

        while ( right != NULL ) { // 求右子樹深度
            right = right -> right;
            rightDepth++;
        }

        if ( leftDepth == rightDepth ) {
            return pow( 2, leftDepth) - 1;
        }

        return countNodes( root -> left ) + countNodes( root -> right ) + 1;
    }
};