统计完全二叉树的节点数

统计完全二叉树的节点数

题目

给定一棵完全二叉树，求节点个数，要求时间复杂度小于O(n)

思路

如图根为A的完全二叉树，先求出A的极左节点H的层数4，
然后找到A的右孩子C，然后找到C的极左节点L，L的层数等于整棵树的层数4，
所以可以肯定A的左子树必定是一棵满二叉树，节点数2**(4-1) - 1 = 7，加上A有8个节点，
而A的右子树节点数可以递归去求解。

如果A的右孩子C的极左节点F没有到达最后一层，则说明A的右子树必定是满二叉树，节点数为2**(3 - 1) - 1，连上根节点，共4个节点，而左子树可以递归求解。

实现

class TreeNode():
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None


def node_num_of_cbt(root):
    def most_left_level(node, level):
        while node:
            level += 1
            node = node.left
        return level - 1

    def node_num(node, level):
        if level == height:
            return 1

        if most_left_level(node.right, level+1) == height:
            return (1 << (height - level)) + node_num(node.right, level+1)
        else:
            return (1 << (height - level - 1)) + node_num(node.left, level+1)

    if root is None:
        return 0

    height = most_left_level(root, 0)
    return node_num(root, 0)

测试

import unittest
import node_num_of_complete_binary_tree as cbt
import random


class NodeNumOfCompleteBinaryTreeCase(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

    def make_cbt(self, count):
        if count == 0:
            return None

        nodes = [cbt.TreeNode(i) for i in range(count)]
        for i in range(count):
            left = 2*i + 1
            right = 2*i + 2
            if left < count:
                nodes[i].left = nodes[left]
            if right < count:
                nodes[i].right = nodes[right]

        return nodes[0]

    def test_sequence(self):
        for i in range(10):
            root = self.make_cbt(i)
            cnt = cbt.node_num_of_cbt(root)
            self.assertTrue(cnt == i)

    def test_random(self):
        for i in range(10):
            count = random.randint(0, 100)
            root = self.make_cbt(count)
            cnt = cbt.node_num_of_cbt(root)
            self.assertTrue(cnt == count)

结果

➜  24_node_nuim_of_complete_binary_tree python -m unittest test_node_num_of_complete_binary_tree
..
----------------------------------------------------------------------
Ran 2 tests in 0.001s

OK