Python | Leetcode Python题解之第230题二叉搜索树中第K小的元素

题目：

题解：

class AVL:
    """平衡二叉搜索树（AVL树）：允许重复值"""

    class Node:
        """平衡二叉搜索树结点"""
        __slots__ = ("val", "parent", "left", "right", "size", "height")

        def __init__(self, val, parent=None, left=None, right=None):
            self.val = val
            self.parent = parent
            self.left = left
            self.right = right
            self.height = 0  # 结点高度：以node为根节点的子树的高度（高度定义：叶结点的高度是0）
            self.size = 1  # 结点元素数：以node为根节点的子树的节点总数

    def __init__(self, vals):
        self.root = self._build(vals, 0, len(vals) - 1, None) if vals else None

    def _build(self, vals, l, r, parent):
        """根据vals[l:r]构造平衡二叉搜索树 -> 返回根结点"""
        m = (l + r) // 2
        node = self.Node(vals[m], parent=parent)
        if l <= m - 1:
            node.left = self._build(vals, l, m - 1, parent=node)
        if m + 1 <= r:
            node.right = self._build(vals, m + 1, r, parent=node)
        self._recompute(node)
        return node

    def kth_smallest(self, k: int) -> int:
        """返回二叉搜索树中第k小的元素"""
        node = self.root
        while node:
            left = self._get_size(node.left)
            if left < k - 1:
                node = node.right
                k -= left + 1
            elif left == k - 1:
                return node.val
            else:
                node = node.left

    def insert(self, v):
        """插入值为v的新结点"""
        if self.root is None:
            self.root = self.Node(v)
        else:
            # 计算新结点的添加位置
            node = self._subtree_search(self.root, v)
            is_add_left = (v <= node.val)  # 是否将新结点添加到node的左子结点
            if node.val == v:  # 如果值为v的结点已存在
                if node.left:  # 值为v的结点存在左子结点，则添加到其左子树的最右侧
                    node = self._subtree_last(node.left)
                    is_add_left = False
                else:  # 值为v的结点不存在左子结点，则添加到其左子结点
                    is_add_left = True

            # 添加新结点
            leaf = self.Node(v, parent=node)
            if is_add_left:
                node.left = leaf
            else:
                node.right = leaf

            self._rebalance(leaf)

    def delete(self, v) -> bool:
        """删除值为v的结点 -> 返回是否成功删除结点"""
        if self.root is None:
            return False

        node = self._subtree_search(self.root, v)
        if node.val != v:  # 没有找到需要删除的结点
            return False

        # 处理当前结点既有左子树也有右子树的情况
        # 若左子树比右子树高度低，则将当前结点替换为右子树最左侧的结点，并移除右子树最左侧的结点
        # 若右子树比左子树高度低，则将当前结点替换为左子树最右侧的结点，并移除左子树最右侧的结点
        if node.left and node.right:
            if node.left.height <= node.right.height:
                replacement = self._subtree_first(node.right)
            else:
                replacement = self._subtree_last(node.left)
            node.val = replacement.val
            node = replacement

        parent = node.parent
        self._delete(node)
        self._rebalance(parent)
        return True

    def _delete(self, node):
        """删除结点p并用它的子结点代替它，结点p至多只能有1个子结点"""
        if node.left and node.right:
            raise ValueError('node has two children')
        child = node.left if node.left else node.right
        if child is not None:
            child.parent = node.parent
        if node is self.root:
            self.root = child
        else:
            parent = node.parent
            if node is parent.left:
                parent.left = child
            else:
                parent.right = child
        node.parent = node

    def _subtree_search(self, node, v):
        """在以node为根结点的子树中搜索值为v的结点，如果没有值为v的结点，则返回值为v的结点应该在的位置的父结点"""
        if node.val < v and node.right is not None:
            return self._subtree_search(node.right, v)
        elif node.val > v and node.left is not None:
            return self._subtree_search(node.left, v)
        else:
            return node

    def _recompute(self, node):
        """重新计算node结点的高度和元素数"""
        node.height = 1 + max(self._get_height(node.left), self._get_height(node.right))
        node.size = 1 + self._get_size(node.left) + self._get_size(node.right)

    def _rebalance(self, node):
        """从node结点开始（含node结点）逐个向上重新平衡二叉树，并更新结点高度和元素数"""
        while node is not None:
            old_height, old_size = node.height, node.size
            if not self._is_balanced(node):
                node = self._restructure(self._tall_grandchild(node))
                self._recompute(node.left)
                self._recompute(node.right)
            self._recompute(node)
            if node.height == old_height and node.size == old_size:
                node = None  # 如果结点高度和元素数都没有变化则不需要再继续向上调整
            else:
                node = node.parent

    def _is_balanced(self, node):
        """判断node结点是否平衡"""
        return abs(self._get_height(node.left) - self._get_height(node.right)) <= 1

    def _tall_child(self, node):
        """获取node结点更高的子树"""
        if self._get_height(node.left) > self._get_height(node.right):
            return node.left
        else:
            return node.right

    def _tall_grandchild(self, node):
        """获取node结点更高的子树中的更高的子树"""
        child = self._tall_child(node)
        return self._tall_child(child)

    @staticmethod
    def _relink(parent, child, is_left):
        """重新连接父结点和子结点（子结点允许为空）"""
        if is_left:
            parent.left = child
        else:
            parent.right = child
        if child is not None:
            child.parent = parent

    def _rotate(self, node):
        """旋转操作"""
        parent = node.parent
        grandparent = parent.parent
        if grandparent is None:
            self.root = node
            node.parent = None
        else:
            self._relink(grandparent, node, parent == grandparent.left)

        if node == parent.left:
            self._relink(parent, node.right, True)
            self._relink(node, parent, False)
        else:
            self._relink(parent, node.left, False)
            self._relink(node, parent, True)

    def _restructure(self, node):
        """trinode操作"""
        parent = node.parent
        grandparent = parent.parent

        if (node == parent.right) == (parent == grandparent.right):  # 处理需要一次旋转的情况
            self._rotate(parent)
            return parent
        else:  # 处理需要两次旋转的情况：第1次旋转后即成为需要一次旋转的情况
            self._rotate(node)
            self._rotate(node)
            return node

    @staticmethod
    def _subtree_first(node):
        """返回以node为根结点的子树的第1个元素"""
        while node.left is not None:
            node = node.left
        return node

    @staticmethod
    def _subtree_last(node):
        """返回以node为根结点的子树的最后1个元素"""
        while node.right is not None:
            node = node.right
        return node

    @staticmethod
    def _get_height(node) -> int:
        """获取以node为根结点的子树的高度"""
        return node.height if node is not None else 0

    @staticmethod
    def _get_size(node) -> int:
        """获取以node为根结点的子树的结点数"""
        return node.size if node is not None else 0


class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        def inorder(node):
            if node.left:
                inorder(node.left)
            inorder_lst.append(node.val)
            if node.right:
                inorder(node.right)

        # 中序遍历生成数值列表
        inorder_lst = []
        inorder(root)

        # 构造平衡二叉搜索树
        avl = AVL(inorder_lst)

        # 模拟1000次插入和删除操作
        random_nums = [random.randint(0, 10001) for _ in range(1000)]
        for num in random_nums:
            avl.insert(num)
        random.shuffle(random_nums)  # 列表乱序
        for num in random_nums:
            avl.delete(num)

        return avl.kth_smallest(k)