使用O(1)额外空间的BST中的第K个最小元素

简介

在二叉搜索树（Binary Search Tree，BST）中，找到第K个最小的元素是一个常见的问题。通常，我们可以通过中序遍历整个BST来得到一个有序的元素列表，然后直接返回第K个元素。但是，这种方法需要使用O(n)的额外空间，其中n是BST中的节点数量。本文介绍了一种使用O(1)额外空间的方法来解决这个问题。

思路

我们可以使用Morris Traversal算法来遍历BST，该算法能够使用O(1)额外空间在不破坏原始树的结构的情况下，通过修改节点的右孩子指针来实现遍历。对于第K个最小的元素，我们可以进行类似的遍历，但是需要在遍历过程中进行一些额外的处理。

我们首先初始化当前节点为根节点，然后开始遍历BST。

1. 如果当前节点的左孩子为空，我们直接访问当前节点，并将当前节点更新为其右孩子。
2. 如果当前节点的左孩子不为空，我们将当前节点的左子树中最右侧的节点设置为当前节点的前驱（predecessor），这个前驱节点是当前节点的直接前一个节点。
   • 如果前驱节点的右孩子为空，我们将它的右孩子指向当前节点，并将当前节点更新为其左孩子。这个步骤是为了在遍历完前驱节点的左子树后，能够返回当前节点。
   • 如果前驱节点的右孩子指向当前节点，我们将它的右孩子重置为空，并访问当前节点。然后将当前节点更新为其右孩子。
3. 重复上述步骤1和步骤2，直到当前节点为空。

在这个算法中，我们将遍历过程和查找第K个最小元素的操作结合在一起，使得我们只需要O(1)的额外空间。

代码实现

以下是用Python实现上述思路的示例代码，其中TreeNode是BST的节点类型：

def find_kth_smallest(root, k):
    count = 0
    curr = root
    kth_smallest = None
    
    while curr:
        if curr.left is None:
            count += 1
            if count == k:
                kth_smallest = curr.val
            curr = curr.right
        else:
            predecessor = curr.left
            while predecessor.right and predecessor.right != curr:
                predecessor = predecessor.right

            if predecessor.right is None:
                predecessor.right = curr
                curr = curr.left
            else:
                predecessor.right = None
                count += 1
                if count == k:
                    kth_smallest = curr.val
                curr = curr.right
    
    return kth_smallest

复杂度分析

• 时间复杂度：该算法的时间复杂度为O(h+k)，其中h是BST的高度。在最坏的情况下，需要遍历整个树直到找到第K个最小元素，这个时间复杂度是线性的。
• 空间复杂度：由于算法只使用了有限的额外空间，其空间复杂度为O(1)。

总结

通过使用Morris Traversal算法，我们可以在不使用额外空间的情况下，在BST中找到第K个最小的元素。这种方法的关键是通过修改节点的右孩子指针来实现遍历。这种实现方式在面试中是一个很好的展示对二叉树遍历算法和BST特性理解的机会。