题目介绍

这是 GATE-CS-2017 套装1 中的第39道题目，题目要求实现一个基于树的数据结构，支持以下几种操作：

• 插入节点
• 删除节点
• 查找最小值
• 查找最大值
• 查找指定节点的前驱
• 查找指定节点的后继
• 中序遍历

树节点定义

首先我们需要定义树的节点结构：

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

树的节点包含一个值属性 val，以及左右子树指针 left 和 right。

树的实现

然后我们可以实现一个二叉搜索树的类，包含插入、删除、查找、以及遍历等方法：

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

    def insert(self, val):
        def insert_node(root, val):
            if root is None:
                return TreeNode(val)
            elif val < root.val:
                root.left = insert_node(root.left, val)
            elif val > root.val:
                root.right = insert_node(root.right, val)
            return root

        self.root = insert_node(self.root, val)

    def delete(self, val):
        def delete_node(root, val):
            if root is None:
                return root
            elif val < root.val:
                root.left = delete_node(root.left, val)
            elif val > root.val:
                root.right = delete_node(root.right, val)
            else:
                if root.left is None:
                    return root.right
                elif root.right is None:
                    return root.left
                else:
                    min_node = root.right
                    while min_node.left is not None:
                        min_node = min_node.left
                    root.val = min_node.val
                    root.right = delete_node(root.right, min_node.val)
            return root

        self.root = delete_node(self.root, val)

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

        return search_node(self.root, val)

    def find_min(self):
        def find_min_node(root):
            if root is None:
                return None
            while root.left is not None:
                root = root.left
            return root.val

        return find_min_node(self.root)

    def find_max(self):
        def find_max_node(root):
            if root is None:
                return None
            while root.right is not None:
                root = root.right
            return root.val

        return find_max_node(self.root)

    def predecessor(self, val):
        def find_node(root, val):
            if root is None:
                return None
            elif val == root.val:
                return root
            elif val < root.val:
                return find_node(root.left, val)
            else:
                return find_node(root.right, val)

        node = find_node(self.root, val)
        if node is None:
            return None
        elif node.left is not None:
            cur = node.left
            while cur.right is not None:
                cur = cur.right
            return cur.val
        else:
            cur = self.root
            pre = None
            while cur is not None:
                if cur.val == val:
                    break
                elif cur.val < val:
                    pre = cur
                    cur = cur.right
                else:
                    cur = cur.left
            if pre is None:
                return None
            else:
                return pre.val

    def successor(self, val):
        def find_node(root, val):
            if root is None:
                return None
            elif val == root.val:
                return root
            elif val < root.val:
                return find_node(root.left, val)
            else:
                return find_node(root.right, val)

        node = find_node(self.root, val)
        if node is None:
            return None
        elif node.right is not None:
            cur = node.right
            while cur.left is not None:
                cur = cur.left
            return cur.val
        else:
            cur = self.root
            suc = None
            while cur is not None:
                if cur.val == val:
                    break
                elif cur.val > val:
                    suc = cur
                    cur = cur.left
                else:
                    cur = cur.right
            if suc is None:
                return None
            else:
                return suc.val

    def inorder(self):
        def inorder_node(root, res):
            if root is not None:
                inorder_node(root.left, res)
                res.append(root.val)
                inorder_node(root.right, res)
            return res

        return inorder_node(self.root, [])

测试代码

最后我们可以写一些测试代码来测试我们的二叉搜索树的实现：

bst = BST()

values = [5, 2, 8, 1, 4, 3, 6, 9, 7]
for val in values:
    bst.insert(val)

assert bst.find_min() == 1
assert bst.find_max() == 9
assert bst.successor(3) == 4
assert bst.predecessor(6) == 5
assert bst.search(1) == True
assert bst.search(10) == False
assert bst.inorder() == [1, 2, 3, 4, 5, 6, 7, 8, 9]

bst.delete(4)

assert bst.successor(3) == 5
assert bst.predecessor(6) == 5
assert bst.search(4) == False
assert bst.inorder() == [1, 2, 3, 5, 6, 7, 8, 9]

这些测试代码用于测试二叉搜索树的基本操作是否正确。