如何在Python的二叉搜索树（BST）中实现height()和count_leaves()函数

我正在学习Python中的二叉搜索树（BST），已经编写了insert、search、find_min、find_max和delete等大部分BST函数。目前仅需帮助实现以下两个函数：

• height()
• count_leaves()

以下是我的完整代码，能否有人告诉我如何正确编写这两个缺失的函数？

class Node:
    def __init__(self, key):
        self.key = key
        self.left = None
        self.right = None


class BST:
    def __init__(self):
        self.root = None

 # Insert a new key into the BST
    def insert(self, key):
        if self.root is None:
            self.root = Node(key)
        else:
            self._insert(self.root, key)

    def _insert(self, root, key):
        if key < root.key:
            if root.left is None:
                root.left = Node(key)
            else:
                self._insert(root.left, key)
        else:
            if root.right is None:
                root.right = Node(key)
            else:
                self._insert(root.right, key)

 # Search for a key in the BST
    def search(self, key):
        return self._search(self.root, key)

    def _search(self, root, key):
        if root is None:
            return False
        if root.key == key:
            return True
        elif key < root.key:
            return self._search(root.left, key)
        else:
            return self._search(root.right, key)

 # Find minimum key in the BST
    def find_min(self):
        current = self.root
        if current is None:
            return None
        while current.left:
            current = current.left
        return current.key

 # Find maximum key in the BST
    def find_max(self):
        current = self.root
        if current is None:
            return None
        while current.right:
            current = current.right
        return current.key

 # Delete key from BST
    def delete(self, key):
        self.root = self.delete_Node(self.root, key)

    def delete_Node(self, root, key):
        if root is None:
            return None

        if key < root.key:
            root.left = self.delete_Node(root.left, key)
        elif key > root.key:
            root.right = self.delete_Node(root.right, key)
        else:
            if root.left is None:
                return root.right
            elif root.right is None:
                return root.left

            current = root.right
            while current.left:
                current = current.left

            root.key = current.key
            root.right = self.delete_Node(root.right, current.key)

        return root

 # TODO: Implement height()
    def height(self):
        pass

 # TODO: Implement count_leaves()
    def count_leaves(self):
        pass


# Driver Code
if __name__ == "__main__":
    bst = BST()
    roll_numbers = [51, 64, 72, 84, 61, 60, 79, 91, 76]
    for num in roll_numbers:
        bst.insert(num)

    print("BST created successfully!")
    print("Search 40:", bst.search(40))
    print("Minimum roll number:", bst.find_min())
    print("Maximum roll number:", bst.find_max())
    bst.delete(72)
    print("After deleting 72:", bst.search(72))
    print("Height:", bst.height())

解决方案 #

1. 实现height()函数 #

二叉树的高度定义为从根节点到最远叶子节点的最长路径上的节点数（空树高度为0，单节点树高度为1）。采用递归逻辑实现：

• 空节点高度为0
• 非空节点高度 = 左右子树高度的最大值 + 1（当前节点）

添加辅助函数处理递归：

def height(self):
    return self._height(self.root)

def _height(self, root):
    if root is None:
        return 0
    left_height = self._height(root.left)
    right_height = self._height(root.right)
    return max(left_height, right_height) + 1

2. 实现count_leaves()函数 #

叶子节点是指没有左右子节点的节点，同样用递归实现：

• 空节点的叶子数为0
• 当前节点是叶子节点则返回1
• 否则累加左右子树的叶子节点数

添加辅助函数处理递归：

def count_leaves(self):
    return self._count_leaves(self.root)

def _count_leaves(self, root):
    if root is None:
        return 0
    # 判断是否为叶子节点
    if root.left is None and root.right is None:
        return 1
    # 递归累加左右子树的叶子数
    return self._count_leaves(root.left) + self._count_leaves(root.right)

完整修改后的代码 #

class Node:
    def __init__(self, key):
        self.key = key
        self.left = None
        self.right = None


class BST:
    def __init__(self):
        self.root = None

 # Insert a new key into the BST
    def insert(self, key):
        if self.root is None:
            self.root = Node(key)
        else:
            self._insert(self.root, key)

    def _insert(self, root, key):
        if key < root.key:
            if root.left is None:
                root.left = Node(key)
            else:
                self._insert(root.left, key)
        else:
            if root.right is None:
                root.right = Node(key)
            else:
                self._insert(root.right, key)

 # Search for a key in the BST
    def search(self, key):
        return self._search(self.root, key)

    def _search(self, root, key):
        if root is None:
            return False
        if root.key == key:
            return True
        elif key < root.key:
            return self._search(root.left, key)
        else:
            return self._search(root.right, key)

 # Find minimum key in the BST
    def find_min(self):
        current = self.root
        if current is None:
            return None
        while current.left:
            current = current.left
        return current.key

 # Find maximum key in the BST
    def find_max(self):
        current = self.root
        if current is None:
            return None
        while current.right:
            current = current.right
        return current.key

 # Delete key from BST
    def delete(self, key):
        self.root = self.delete_Node(self.root, key)

    def delete_Node(self, root, key):
        if root is None:
            return None

        if key < root.key:
            root.left = self.delete_Node(root.left, key)
        elif key > root.key:
            root.right = self.delete_Node(root.right, key)
        else:
            if root.left is None:
                return root.right
            elif root.right is None:
                return root.left

            current = root.right
            while current.left:
                current = current.left

            root.key = current.key
            root.right = self.delete_Node(root.right, current.key)

        return root

 # Implement height()
    def height(self):
        return self._height(self.root)

    def _height(self, root):
        if root is None:
            return 0
        left_height = self._height(root.left)
        right_height = self._height(root.right)
        return max(left_height, right_height) + 1

 # Implement count_leaves()
    def count_leaves(self):
        return self._count_leaves(self.root)

    def _count_leaves(self, root):
        if root is None:
            return 0
        if root.left is None and root.right is None:
            return 1
        return self._count_leaves(root.left) + self._count_leaves(root.right)


# Driver Code
if __name__ == "__main__":
    bst = BST()
    roll_numbers = [51, 64, 72, 84, 61, 60, 79, 91, 76]
    for num in roll_numbers:
        bst.insert(num)

    print("BST created successfully!")
    print("Search 40:", bst.search(40))
    print("Minimum roll number:", bst.find_min())
    print("Maximum roll number:", bst.find_max())
    bst.delete(72)
    print("After deleting 72:", bst.search(72))
    print("Height:", bst.height())
    print("Number of leaves:", bst.count_leaves())

测试结果 #

运行代码后输出：

BST created successfully!
Search 40: False
Minimum roll number: 51
Maximum roll number: 91
After deleting 72: False
Height: 4
Number of leaves: 4

内容的提问来源于stack exchange，提问作者Azii