LeetCode 222. Count Complete Tree Nodes

Post by ailswan Nov.05, 2023

222. Count Complete Tree Nodes

Problem Statement

link: https://leetcode.com/problems/count-complete-tree-nodes/ https://leetcode.cn/problems/count-complete-tree-nodes/

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:

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

Input: root = [] Output: 0

Input: root = [1] Output: 1

Solution Approach

The approach is to find the height of the leftmost path and traverse the tree using binary search. At each step, determine if the last node is on the left or right side based on the right subtree’s height, adding the appropriate nodes to the result.

Algorithm

  1. Calculate the leftmost path’s height with a recursive function.
  2. Initialize the result and tree height.
  3. Traverse the tree, adjusting the result based on the right subtree’s height to count nodes efficiently.

Implement

    class Solution:
      def countNodes(self, root: Optional[TreeNode]) -> int:
          def getHeight(node):
              return 1 + getHeight(node.left) if node else -1
  
          res = 0
          h = getHeight(root)
          if h < 0:
              return 0
          while root:
              if getHeight(root.right) == h - 1:
                  res += 2 ** h
                  root = root.right
              else:
                  res += 2 **(h - 1)
                  root = root.left
              h -= 1
          return res