[16]填充每个节点的下一个右侧节点指针和二叉搜索树中第K小的元素

内容来自leetcode

1.填充每个节点的下一个右侧节点指针

题目要求

给定一个 完美二叉树 ，其所有叶子节点都在同一层，每个父节点都有两个子节点。二叉树定义如下：

struct Node {
  int val;
  Node *left;
  Node *right;
  Node *next;
}
填充它的每个 next 指针，让这个指针指向其下一个右侧节点。如果找不到下一个右侧节点，则将 next 指针设置为 NULL。

初始状态下，所有 next 指针都被设置为 NULL。

示例 1：

输入：root = [1,2,3,4,5,6,7]
输出：[1,#,2,3,#,4,5,6,7,#]

解释：给定二叉树如图 A 所示，你的函数应该填充它的每个 next 指针，以指向其下一个右侧节点，如图 B 所示。序列化的输出按层序遍历排列，同一层节点由 next 指针连接，'#' 标志着每一层的结束。

思路

迭代，从root节点开始遍历每个节点的左右儿子，并按先后顺序入栈。每次循坏开始时栈中节点一定都属于一层，并且包含一层中所有节点。遍历完栈中现有节点，得到每个节点的next并将下一层节点入栈。时间复杂度和空间复杂度均为O(N)

"""
# Definition for a Node.
class Node:
    def __init__(self, val: int = 0, left: 'Node' = None, right: 'Node' = None, next: 'Node' = None):
        self.val = val
        self.left = left
        self.right = right
        self.next = next
"""

class Solution:
    def connect(self, root: 'Optional[Node]') -> 'Optional[Node]':
        if not root:
            return root
        #得到初始栈
        stk = [root]
        while stk:
            n = len(stk)
            #以栈中现有节点数为循环次数
            for i in range(n):
                #将第一个节点出栈
                node = stk.pop(0)
                #将节点左右节点入栈（如果有）
                if node.left:
                    stk.append(node.left)
                if node.right:
                    stk.append(node.right)
                #判断新栈首节点是否属于这轮循环内节点
                if i + 1 < n:
                    node.next = stk[0]
        return root

第二种解法基本思路是，一个节点的右儿子一定是其左儿子的next节点，而其右儿子的next节点一定是节点右兄弟的左儿子。也就是说，可以通过遍历第N层，来建立N+1层节点next关系。这种方法将空间复杂度降低为了O(1)。

自己写的

class Solution:
    def connect(self, root: 'Optional[Node]') -> 'Optional[Node]':
        head = root
        first = root
        while root:
            rootLeft = first.left
            if root.left:
                root.left.next = root.right
            if root.next and root.right:
                root.right.next = root.next.left
                root = root.next
            else:
                root = rootLeft
                first = root
        return head

官方解法

class Solution:
    def connect(self, root: 'Node') -> 'Node':
        
        if not root:
            return root
        
        # 从根节点开始
        leftmost = root
        
        while leftmost.left:
            
            # 遍历这一层节点组织成的链表，为下一层的节点更新 next 指针
            head = leftmost
            while head:
                
                # CONNECTION 1
                head.left.next = head.right
                
                # CONNECTION 2
                if head.next:
                    head.right.next = head.next.left
                
                # 指针向后移动
                head = head.next
            
            # 去下一层的最左的节点
            leftmost = leftmost.left
        
        return root

2.二叉搜索树中第K小的元素

题目要求

给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 个最小元素（从 1 开始计数）。

示例 1：

输入：root = [3,1,4,null,2], k = 1
输出：1

进阶：如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化算法？

思路

最为直接的办法就是把树遍历一遍，然后将遍历得到的列表排序，输出。。。

但是题目给出的是二叉搜索树，节点排布是有规律的，中序遍历得到的就是排序过的列表。所以就有了第一种解法。

class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        def inorder(root):
            if root == None:
                return []
            else:
                result = []

                result = result + inorder(root.left)               
                result.append(root.val)
                result = result + inorder(root.right)
            return result
        treeList = inorder(root)
        return treeList[k-1]

官方给出的第二种解法是记录子树的结点数

因此，我们可以记录下以每个结点为根结点的子树的结点数，并在查找第 k 小的值时，使用如下方法搜索：

令 node等于根结点，开始搜索。

对当前结点 node进行如下操作：

如果 node的左子树的结点数 left小于 k−1，则第 k 小的元素一定在 node 的右子树中，令 node 等于其的右子结点，k 等于 k−left−1，并继续搜索；
如果 node 的左子树的结点数 left 等于 k−1，则第 k 小的元素即为 node，结束搜索并返回 node即可；
如果 node 的左子树的结点数 left 大于 k−1，则第 k 小的元素一定在 node 的左子树中，令 node等于其左子结点，并继续搜索。
在实现中，我们既可以将以每个结点为根结点的子树的结点数存储在结点中，也可以将其记录在哈希表中。

class nodeNum:
    def __init__(self,root):
        self.root = root

    def countNodeNum(self,root):
    #取得当前节点的子节点数（包括自身的）
        if not root:
            return 0
        stk = [root]
        count = 1
        while stk:
            node = stk.pop()
            if node.left:
                count += 1
                stk.append(node.left)
            if node.right:
                count += 1
                stk.append(node.right)
        return count
    def findKthSmallest(self,node,k):
    #查找第K小的节点位置
        #每次需要找的时候去获取节点左子树节点数
        leftNodeNum = self.countNodeNum(node.left)
        if leftNodeNum < k-1:
            node = node.right
            k = k - leftNodeNum - 1
            node = self.findKthSmallest(node,k)
        elif leftNodeNum == k-1:
            return node.val
        else:
            node = node.left
            node = self.findKthSmallest(node,k)
        return node
class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        bst = nodeNum(root)
        return bst.findKthSmallest(root,k)

官方是用哈希表来存储每个节点的子树的节点数的～

class MyBst:
    def __init__(self, root: TreeNode):
        self.root = root

        # 统计以每个结点为根结点的子树的结点数，并存储在哈希表中
        self._node_num = {}
        self._count_node_num(root)

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

    def _count_node_num(self, node) -> int:
        """统计以node为根结点的子树的结点数"""
        if not node:
            return 0
        self._node_num[node] = 1 + self._count_node_num(node.left) + self._count_node_num(node.right)
        return self._node_num[node]

    def _get_node_num(self, node) -> int:
        """获取以node为根结点的子树的结点数"""
        return self._node_num[node] if node is not None else 0


class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        bst = MyBst(root)
        return bst.kth_smallest(k)

第三种解法是用平衡二叉搜索树

平衡二叉搜索树（AVL树） 的知识。平衡二叉搜索树具有如下性质：

平衡二叉搜索树中每个结点的左子树和右子树的高度最多相差 1；

平衡二叉搜索树的子树也是平衡二叉搜索树；

一棵存有 n 个结点的平衡二叉搜索树的高度是 O(log⁡n)。

我们注意到在方法二中搜索二叉搜索树的时间复杂度为 O(H)，其中 HH是树的高度；当树是平衡树时，时间复杂度取得最小值 O(log⁡N)。因此，我们在记录子树的结点数的基础上，将二叉搜索树转换为平衡二叉搜索树，并在插入和删除操作中维护它的平衡状态。

这里给出代码供参考

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)