BST 中序有序,迭代栈访问到第 k 个就停止。
OJ: leetcodecn
题目 ID: kth-smallest-element-in-a-bst
难度:普及+/提高
标签:BST中序栈cpppython
日期: 2026-07-29 13:10
题意
返回 BST 中第 k 小的元素。
思路
BST 中序遍历 = 升序序列。用迭代栈遍历到第 k 个即返回。
代码
cpp
#include <bits/stdc++.h>
using namespace std;
struct TreeNode {
int val;
TreeNode *left, *right;
TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
};
class Solution {
public:
int kthSmallest(TreeNode *root, int k) {
stack<TreeNode *> st;
auto cur = root;
while (cur || !st.empty()) {
while (cur) {
st.push(cur);
cur = cur->left;
}
cur = st.top();
st.pop();
if (--k == 0)
return cur->val;
cur = cur->right;
}
return -1;
}
};
TreeNode *build(istream &in, int n) {
if (!n)
return nullptr;
vector<TreeNode *> nodes(n);
queue<TreeNode *> q;
for (int i = 0, v; i < n; i++) {
in >> v;
if (v != -1)
nodes[i] = new TreeNode(v);
}
TreeNode *root = nodes[0];
if (root)
q.push(root);
int idx = 1;
while (!q.empty() && idx < n) {
auto cur = q.front();
q.pop();
if (idx < n) {
cur->left = nodes[idx];
if (nodes[idx])
q.push(nodes[idx]);
idx++;
}
if (idx < n) {
cur->right = nodes[idx];
if (nodes[idx])
q.push(nodes[idx]);
idx++;
}
}
return root;
}
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n, k;
cin >> n >> k;
auto r = build(cin, n);
cout << Solution().kthSmallest(r, k) << '\n';
return 0;
}python
#!/usr/bin/env python3
from collections import deque
from typing import Optional
class TreeNode:
def __init__(self, x):
self.val = x
self.left = self.right = None
class Solution:
def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
st, cur = [], root
while cur or st:
while cur:
st.append(cur)
cur = cur.left
cur = st.pop()
k -= 1
if k == 0:
return cur.val
cur = cur.right
return -1
def build(arr):
if not arr:
return None
nodes = [TreeNode(v) if v != -1 else None for v in arr]
q = deque([nodes[0]]) if nodes[0] else deque()
idx = 1
while q and idx < len(arr):
cur = q.popleft()
if idx < len(arr):
cur.left = nodes[idx]
if nodes[idx]:
q.append(nodes[idx])
idx += 1
if idx < len(arr):
cur.right = nodes[idx]
if nodes[idx]:
q.append(nodes[idx])
idx += 1
return nodes[0]
def main():
n, k = map(int, input().split())
a = list(map(int, input().split()))
print(Solution().kthSmallest(build(a), k))
if __name__ == "__main__":
main()复杂度
时间 O(H + k),空间 O(H)。
总结
BST 的中序性质让第 k 小问题变成"控制中序遍历的终止时机"。