題目出處

222. Count Complete Tree Nodes

難度

medium

個人範例程式碼

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def countNodes(self, root: Optional[TreeNode]) -> int:
        if not root:
            return 0

        left_depth = self.get_depth(root.left)
        right_depth = self.get_depth(root.right)

        if left_depth == right_depth: # (left completed)
            return 1 + (2**left_depth - 1) + self.countNodes(root.right)
        else: # left_depth > right_depth (right completed)
            return 1 + (2**right_depth - 1) + self.countNodes(root.left)

    def get_depth(self, root):
        if not root:
            return 0

        return 1 + self.get_depth(root.left)

算法說明

拆成兩路進行，兩路我們都直接找「最左下」，比較兩者的深度，此時會有兩種情況

• 當左 右，表示左邊一定是 completed binary tree
• 當左 > 右，表示右邊一定是 completed binary tree

你可能想問，那「左 < 右」?
不可能，你再想想 completed Tree 的定義XD

因此，我們簡略圖示一下，當「左 右」，左邊一定是 completed

當「左 < 右」，右邊一定是 completed

我們可以搭配 「2**depth -1」快速計算出 completed tree 的總 node 數。

input handling

處理沒有 root 的情況，return 0

Boundary conditions

用「not root」去控制 traverse tree 的結束。

Reference

• 统计完全树节点数 · Count Complete Tree Nodes