Python——计算完全二叉树的节点个数

完全二叉树的定义

在完全二叉树中，除了最底层结点可能没填满外，其余每层结点数都达到最大值，并且最下面一层的结点都集中在该层最左边的若干位置。若最底层为第 h 层，则该层包含 1~ 2^h 个结点

预期结果

输入： [1,2,3,4,5,6]
输出： 6

递归法

class Solution:
    def countNodes(self, root: TreeNode) -> int:
        left = right = 0
        tmp1 = tmp2 =  root
        # 计算左子树深度
        while tmp1:
            left += 1
            tmp1 = tmp1.left
        # 计算右子树深度 
        while tmp2:
            right += 1
            tmp2 = tmp2.right
        if left == right:
            return 2 ** left - 1
        else:
            return 1 + self.countNodes(root.left) + self.countNodes(root.right)

迭代法

class Solution:
    def getDepth(self, root):  # 计算当前树的深度
        depth = 0
        while root:
            depth += 1
            root = root.left
        return depth
    
    def countNodes(self, root: TreeNode) -> int:
        count = 0
        while root:
            left = self.getDepth(root.left)  # 左子树深度
            right = self.getDepth(root.right)  # 右子树深度
            # 左右子树深度相同，左子树一定是满二叉树，右子树可能为满二叉树
            if left == right:  
                root = root.right
                count += 2 ** left
            # 左右子树深度不同，右子树一定是满二叉树，左子树可能为满二叉树
            else:              
                root = root.left
                count += 2 ** right  
        return count