leetcode地址:二叉搜索树中第K小的元素
给定一个二叉搜索树的根节点 root ,和一个整数 k ,请你设计一个算法查找其中第 k 小的元素(从 1 开始计数)。
示例 1:
输入:root = [3,1,4,null,2], k = 1
输出:1
示例 2:
输入:root = [5,3,6,2,4,null,null,1], k = 3
输出:3
提示:
树中的节点数为 n 。
1 <= k <= n <= 104
0 <= Node.val <= 104
实现思路
在二叉搜索树(BST)中查找第 k 小的元素。二叉搜索树的特点是对于每个节点,左子树中的所有节点值小于该节点值,右子树中的所有节点值大于该节点值。利用这一性质,我们可以通过中序遍历(in-order traversal)来得到一个有序的元素列表,然后查找第 k 小的元素。
中序遍历
中序遍历(in-order traversal)是指先遍历左子树,再访问根节点,最后遍历右子树。对于二叉搜索树,中序遍历会按顺序访问所有节点,得到一个递增的序列。
记录遍历过程
使用一个变量来记录当前访问到的节点数,当这个计数器等于 k 时,当前节点即为第 k 小的元素。
递归或迭代
我们可以使用递归或迭代的方法来实现中序遍历。
递归实现
# 定义二叉树节点类
class TreeNode:
def __init__(self, val=0, left=None, right=None):
self.val = val
self.left = left
self.right = right
# 查找第k小元素的函数
def kthSmallest(root: TreeNode, k: int) -> int:
def inOrderTraversal(node):
if not node:
return []
# 递归遍历左子树
left = inOrderTraversal(node.left)
# 访问当前节点
current = [node.val]
# 递归遍历右子树
right = inOrderTraversal(node.right)
return left + current + right
# 获取中序遍历结果
result = inOrderTraversal(root)
# 返回第k小的元素
return result[k - 1]
# 测试示例
if __name__ == "__main__":
# 创建测试二叉搜索树
# 3
# / \
# 1 4
# \
# 2
root = TreeNode(3)
root.left = TreeNode(1)
root.right = TreeNode(4)
root.left.right = TreeNode(2)
k = 1
print(f"第{k}小的元素: {kthSmallest(root, k)}") # 应该输出1
迭代实现
# 定义二叉树节点类
class TreeNode:
def __init__(self, val=0, left=None, right=None):
self.val = val
self.left = left
self.right = right
# 查找第k小元素的函数
def kthSmallest(root: TreeNode, k: int) -> int:
stack = []
while True:
# 左子树遍历
while root:
stack.append(root)
root = root.left
root = stack.pop()
k -= 1
# 当计数器等于0时,返回当前节点的值
if k == 0:
return root.val
# 右子树遍历
root = root.right
# 测试示例
if __name__ == "__main__":
# 创建测试二叉搜索树
# 3
# / \
# 1 4
# \
# 2
root = TreeNode(3)
root.left = TreeNode(1)
root.right = TreeNode(4)
root.left.right = TreeNode(2)
k = 1
print(f"第{k}小的元素: {kthSmallest(root, k)}") # 应该输出1
Go实现
递归实现
package main
import "fmt"
// TreeNode 定义二叉树节点
type TreeNode struct {
Val int
Left *TreeNode
Right *TreeNode
}
// kthSmallest 递归查找二叉搜索树中的第k小元素
func kthSmallest(root *TreeNode, k int) int {
var inOrder func(node *TreeNode) []int
inOrder = func(node *TreeNode) []int {
if node == nil {
return []int{}
}
left := inOrder(node.Left)
current := []int{node.Val}
right := inOrder(node.Right)
return append(append(left, current...), right...)
}
result := inOrder(root)
return result[k-1]
}
// 测试示例
func main() {
// 创建测试二叉搜索树
// 3
// / \
// 1 4
// \
// 2
root := &TreeNode{Val: 3}
root.Left = &TreeNode{Val: 1}
root.Right = &TreeNode{Val: 4}
root.Left.Right = &TreeNode{Val: 2}
k := 1
fmt.Printf("第%d小的元素: %d\n", k, kthSmallest(root, k)) // 应该输出1
}
迭代实现
package main
import "fmt"
// TreeNode 定义二叉树节点
type TreeNode struct {
Val int
Left *TreeNode
Right *TreeNode
}
// kthSmallest 迭代查找二叉搜索树中的第k小元素
func kthSmallest(root *TreeNode, k int) int {
stack := []*TreeNode{}
current := root
for current != nil || len(stack) > 0 {
for current != nil {
stack = append(stack, current)
current = current.Left
}
current = stack[len(stack)-1]
stack = stack[:len(stack)-1]
k--
if k == 0 {
return current.Val
}
current = current.Right
}
return -1 // 如果树中没有第k小的元素
}
// 测试示例
func main() {
// 创建测试二叉搜索树
// 3
// / \
// 1 4
// \
// 2
root := &TreeNode{Val: 3}
root.Left = &TreeNode{Val: 1}
root.Right = &TreeNode{Val: 4}
root.Left.Right = &TreeNode{Val: 2}
k := 1
fmt.Printf("第%d小的元素: %d\n", k, kthSmallest(root, k)) // 应该输出1
}
Kotlin实现
递归实现
// 定义二叉树节点类
class TreeNode(var `val`: Int) {
var left: TreeNode? = null
var right: TreeNode? = null
}
// 查找第k小元素的函数
fun kthSmallest(root: TreeNode?, k: Int): Int {
// 中序遍历函数
fun inOrderTraversal(node: TreeNode?): List<Int> {
if (node == null) return listOf()
val left = inOrderTraversal(node.left)
val current = listOf(node.`val`)
val right = inOrderTraversal(node.right)
return left + current + right
}
val result = inOrderTraversal(root)
return result[k - 1]
}
// 测试示例
fun main() {
// 创建测试二叉搜索树
// 3
// / \
// 1 4
// \
// 2
val root = TreeNode(3).apply {
left = TreeNode(1).apply {
right = TreeNode(2)
}
right = TreeNode(4)
}
val k = 1
println("第${k}小的元素: ${kthSmallest(root, k)}") // 应该输出1
}
迭代实现
// 定义二叉树节点类
class TreeNode(var `val`: Int) {
var left: TreeNode? = null
var right: TreeNode? = null
}
// 查找第k小元素的函数
fun kthSmallest(root: TreeNode?, k: Int): Int {
val stack = mutableListOf<TreeNode>()
var current = root
var kVar = k
while (current != null || stack.isNotEmpty()) {
while (current != null) {
stack.add(current)
current = current.left
}
current = stack.removeAt(stack.size - 1)
kVar--
if (kVar == 0) {
return current.`val`
}
current = current.right
}
throw IllegalArgumentException("Tree does not contain k elements")
}
// 测试示例
fun main() {
// 创建测试二叉搜索树
// 3
// / \
// 1 4
// \
// 2
val root = TreeNode(3).apply {
left = TreeNode(1).apply {
right = TreeNode(2)
}
right = TreeNode(4)
}
val k = 1
println("第${k}小的元素: ${kthSmallest(root, k)}") // 应该输出1
}