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

题目要求

设计一个算法，用于从给定的二叉搜索树中找到第 k 个最小的元素。这里的二叉搜索树是指一个每个节点的左子树只包含小于当前节点的数，右子树只包含大于当前节点的数的二叉树。算法需要接收两个输入参数：

1. root：二叉搜索树的根节点。
2. k：需要查找的第 k 个最小元素的序号（序号从 1 开始）。

解题思路

由于二叉搜索树（BST）的性质，我们知道对其进行中序遍历可以得到一个升序排列的元素序列。因此，查找第 k 个最小元素的问题可以转化为执行中序遍历直到遍历到第 k 个节点。

以下是解题的步骤：

1. 中序遍历：按照“左-根-右”的顺序递归遍历二叉搜索树。这样可以保证访问节点的顺序是按照节点值的升序。

2. 计数：在遍历过程中，维护一个计数器，记录已经遍历过的节点数量。

3. 查找第 k 个元素：当计数器的值达到 k 时，当前遍历到的节点值即为所求的第 k 个最小元素。

4. 优化：由于只需要找到第 k 个最小元素，一旦找到，就可以停止遍历，不需要遍历整棵树。

5. 边界情况处理：需要考虑树为空或 k 超出树中元素总数的情况，合理处理这些边界情况。

通过上述步骤，我们可以设计出一个高效的算法来解决这个问题，时间复杂度为 O(h + k)，其中 h 是树的高度。

Golang 版本

package main

import (
	"fmt"
)

// TreeNode is the structure for binary tree nodes
type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

// kthSmallest function finds the kth smallest element in a BST
func kthSmallest(root *TreeNode, k int) int {
	count := 0
	result := 0

	// Recursive function to perform inorder traversal
	var inorder func(node *TreeNode)
	inorder = func(node *TreeNode) {
		if node == nil {
			return
		}

		// Traverse the left subtree
		inorder(node.Left)

		// Increment count and check if the current node is the kth smallest
		count++
		if count == k {
			result = node.Val
			return
		}

		// Traverse the right subtree
		inorder(node.Right)
	}

	// Start inorder traversal from the root
	inorder(root)

	return result
}

func main() {
	// Example usage:
	// Construct a binary search tree here and call kthSmallest with its root and k
}

这段代码定义了一个TreeNode结构体，表示二叉树的节点，并实现了kthSmallest函数，该函数通过中序遍历找到二叉搜索树中第 k 小的元素。main函数中可以构造一个二叉搜索树，并调用kthSmallest函数来测试功能。注意，这段代码只提供了核心函数的实现，实际使用时需要构建二叉树并提供root和k的值。

C++版本

#include <iostream>

// Definition for a binary tree node.
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        int count = 0;
        int result = -1;
        inorder(root, k, count, result);
        return result;
    }

private:
    void inorder(TreeNode* node, int k, int& count, int& result) {
        if (node == nullptr || count >= k) {
            return;
        }

        inorder(node->left, k, count, result);

        if (++count == k) {
            result = node->val;
            return;
        }

        inorder(node->right, k, count, result);
    }
};

int main() {
    // Example usage:
    // TreeNode* root = new TreeNode(3);
    // root->left = new TreeNode(1);
    // root->right = new TreeNode(4);
    // root->left->right = new TreeNode(2);
    // Solution solution;
    // int k = 1;
    // std::cout << "Kth smallest element is: " << solution.kthSmallest(root, k) << std::endl;
    // Output should be 1 for the given tree and k value.

    // Remember to delete the created tree nodes to prevent memory leak.
}

这段代码定义了一个TreeNode结构体，表示二叉树的节点，并实现了一个Solution类，其中包含了kthSmallest方法，该方法通过中序遍历找到二叉搜索树中第 k 小的元素。main函数中可以构造一个二叉搜索树，并调用kthSmallest方法来测试功能。请注意，示例用法中的树节点创建后，应当在适当的时候删除以防止内存泄漏。

Python 版本

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

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        # Initialize the count and the result placeholder
        self.k = k
        self.res = None
        self._inorder(root)
        return self.res

    def _inorder(self, node):
        # If node is None or result is already found, return
        if not node or self.res is not None:
            return

        # Traverse the left subtree
        self._inorder(node.left)

        # Decrement k and check if we have reached the kth node
        self.k -= 1
        if self.k == 0:
            self.res = node.val
            return

        # Traverse the right subtree
        self._inorder(node.right)

# Example usage:
# root = TreeNode(3)
# root.left = TreeNode(1)
# root.right = TreeNode(4)
# root.left.right = TreeNode(2)
# solution = Solution()
# k = 1
# print(f"The {k}th smallest element is: {solution.kthSmallest(root, k)}")

这段代码定义了一个TreeNode类，表示二叉树的节点，并实现了一个Solution类，其中包含了kthSmallest方法，该方法通过中序遍历找到二叉搜索树中第 k 小的元素。示例用法展示了如何创建一个二叉搜索树并调用kthSmallest方法来获取第 k 小的元素。

Java 版本

public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { val = x; }
}

class Solution {
    private int count = 0;
    private int result = Integer.MIN_VALUE;

    public int kthSmallest(TreeNode root, int k) {
        inorder(root, k);
        return result;
    }

    private void inorder(TreeNode node, int k) {
        if (node == null) {
            return;
        }
        inorder(node.left, k);
        if (++count == k) {
            result = node.val;
            return;
        }
        inorder(node.right, k);
    }
}

public class Main {
    public static void main(String[] args) {
        // Example usage:
        // TreeNode root = new TreeNode(3);
        // root.left = new TreeNode(1);
        // root.right = new TreeNode(4);
        // root.left.right = new TreeNode(2);
        // Solution solution = new Solution();
        // int k = 1;
        // System.out.println("The " + k + "th smallest element is: " + solution.kthSmallest(root, k));
    }
}

这段代码定义了一个TreeNode类，表示二叉树的节点，并实现了一个Solution类，其中包含了kthSmallest方法，该方法通过中序遍历找到二叉搜索树中第 k 小的元素。Main类中的main方法提供了如何创建一个二叉搜索树并调用kthSmallest方法来获取第 k 小的元素的示例用法。注意，示例用法中的代码被注释掉了，以防止直接运行时出现错误。在实际使用时，应当取消注释并根据需要构造二叉树。

Rust 版本

// Definition for a binary tree node.
#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
    pub val: i32,
    pub left: Option<Rc<RefCell<TreeNode>>>,
    pub right: Option<Rc<RefCell<TreeNode>>>,
}

impl TreeNode {
    #[inline]
    pub fn new(val: i32) -> Self {
        TreeNode {
            val,
            left: None,
            right: None,
        }
    }
}

use std::rc::Rc;
use std::cell::RefCell;

struct Solution;

impl Solution {
    pub fn kth_smallest(root: Option<Rc<RefCell<TreeNode>>>, k: i32) -> i32 {
        let mut stack = Vec::new();
        let mut current = root.clone();
        let mut k = k;

        while current.is_some() || !stack.is_empty() {
            while let Some(node) = current {
                stack.push(node.clone());
                current = node.borrow().left.clone();
            }

            current = stack.pop();
            if let Some(node) = current {
                k -= 1;
                if k == 0 {
                    return node.borrow().val;
                }
                current = node.borrow().right.clone();
            }
        }

        panic!("K is larger than the number of nodes in the tree");
    }
}

fn main() {
    // Example usage:
    // let root = Some(Rc::new(RefCell::new(TreeNode::new(3))));
    // root.as_ref().unwrap().borrow_mut().left = Some(Rc::new(RefCell::new(TreeNode::new(1))));
    // root.as_ref().unwrap().borrow_mut().right = Some(Rc::new(RefCell::new(TreeNode::new(4))));
    // root.as_ref().unwrap().borrow_mut().left.as_ref().unwrap().borrow_mut().right = Some(Rc::new(RefCell::new(TreeNode::new(2))));
    // let k = 1;
    // println!("The {}th smallest element is: {}", k, Solution::kth_smallest(root, k));
}

这段代码使用 Rust 语言定义了一个TreeNode结构体，表示二叉树的节点，并实现了一个Solution结构体，其中包含了kth_smallest函数，该函数通过中序遍历的迭代方式找到二叉搜索树中第 k 小的元素。main函数中提供了如何创建一个二叉搜索树并调用kth_smallest函数来获取第 k 小的元素的示例用法。注意，示例用法中的代码被注释掉了，以防止直接运行时出现错误。在实际使用时，应当取消注释并根据需要构造二叉树。

总结

上述代码段展示了如何在不同编程语言中实现查找二叉搜索树中第 k 小的元素的算法。核心思想是利用二叉搜索树的性质：中序遍历二叉搜索树可以得到一个递增的有序序列。因此，通过中序遍历直到遍历到第 k 个节点，就可以找到第 k 小的元素。

在不同的编程语言中，这一算法的实现有以下共同点：

1. 定义二叉树节点：每种语言都定义了一个TreeNode结构（或类），包含值val，左子节点left和右子节点right。

2. 中序遍历：所有语言都使用了中序遍历的方法来遍历树。中序遍历首先访问左子树，然后访问当前节点，最后访问右子树。

3. 计数器：在遍历过程中，使用一个计数器来记录当前已经遍历过的节点数。当计数器达到 k 时，当前节点的值即为所求的第 k 小的元素。

4. 递归或迭代：根据语言的特性和编程风格的不同，中序遍历可以通过递归或迭代的方式实现。递归实现更为直观，而迭代实现通常需要手动维护一个栈。

5. 返回结果：一旦找到第 k 小的元素，算法就会返回该元素的值。

不同编程语言的实现细节可能有所不同，例如在 Rust 中使用Option<Rc<RefCell<TreeNode>>>来处理节点的所有权和可变性，而在 Java 和 Python 中则直接使用指针或引用。但无论语言如何变化，算法的核心逻辑和步骤是一致的。