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 = valself.parent = parentself.left = leftself.right = rightself.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 Nonedef _build(self, vals, l, r, parent):"""根据vals[l:r]构造平衡二叉搜索树 -> 返回根结点"""m = (l + r) // 2node = 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 nodedef kth_smallest(self, k: int) -> int:"""返回二叉搜索树中第k小的元素"""node = self.rootwhile node:left = self._get_size(node.left)if left < k - 1:node = node.rightk -= left + 1elif left == k - 1:return node.valelse:node = node.leftdef 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 = Falseelse:  # 值为v的结点不存在左子结点，则添加到其左子结点is_add_left = True# 添加新结点leaf = self.Node(v, parent=node)if is_add_left:node.left = leafelse:node.right = leafself._rebalance(leaf)def delete(self, v) -> bool:"""删除值为v的结点 -> 返回是否成功删除结点"""if self.root is None:return Falsenode = 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.valnode = replacementparent = node.parentself._delete(node)self._rebalance(parent)return Truedef _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.rightif child is not None:child.parent = node.parentif node is self.root:self.root = childelse:parent = node.parentif node is parent.left:parent.left = childelse:parent.right = childnode.parent = nodedef _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 nodedef _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.sizeif 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.parentdef _is_balanced(self, node):"""判断node结点是否平衡"""return abs(self._get_height(node.left) - self._get_height(node.right)) <= 1def _tall_child(self, node):"""获取node结点更高的子树"""if self._get_height(node.left) > self._get_height(node.right):return node.leftelse:return node.rightdef _tall_grandchild(self, node):"""获取node结点更高的子树中的更高的子树"""child = self._tall_child(node)return self._tall_child(child)@staticmethoddef _relink(parent, child, is_left):"""重新连接父结点和子结点（子结点允许为空）"""if is_left:parent.left = childelse:parent.right = childif child is not None:child.parent = parentdef _rotate(self, node):"""旋转操作"""parent = node.parentgrandparent = parent.parentif grandparent is None:self.root = nodenode.parent = Noneelse: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.parentgrandparent = parent.parentif (node == parent.right) == (parent == grandparent.right):  # 处理需要一次旋转的情况self._rotate(parent)return parentelse:  # 处理需要两次旋转的情况：第1次旋转后即成为需要一次旋转的情况self._rotate(node)self._rotate(node)return node@staticmethoddef _subtree_first(node):"""返回以node为根结点的子树的第1个元素"""while node.left is not None:node = node.leftreturn node@staticmethoddef _subtree_last(node):"""返回以node为根结点的子树的最后1个元素"""while node.right is not None:node = node.rightreturn node@staticmethoddef _get_height(node) -> int:"""获取以node为根结点的子树的高度"""return node.height if node is not None else 0@staticmethoddef _get_size(node) -> int:"""获取以node为根结点的子树的结点数"""return node.size if node is not None else 0class 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)