今天分享一道王道考研数据结构参考书（P286）中一道算法题:
编写一个递归算法，在一棵有n个结点的、随机建立起来的二叉排序树上查找第k(1≤ k ≤ n）小的元素，并返回指向该结点的指针。要求算法的平均时间复杂度为 O(log2n)。二叉排序树的每个结,点中除data, lchild,rchild 等数据外，增加一个count 成员，保存以该结点为根的子树上的结点个数。
这个题目是找到第 k 小的数，并且是在一棵排序二叉树中，那就是将排序二叉树的节点从小到大排列，从前往后数，第 k 个数呗，是这个意思吧？
要是这么做的话，这个题目就挺简单的，就是中序遍历呗，因为排序二叉树的中序遍历就是一个从小到大的遍历。第 k 小的数就是遍历到的第 k 个数了
准备数据
javascript
const data = [9, 4, 5, 3, 2, 7, 8, 95, 33, 22];
const generator = (data) => {
const insert = (tree, value) => {
if (!tree) return { value, left: null, right: null};
if (value < tree.value) {
tree.left = insert(tree.left, value);
} else {
tree.right = insert(tree.right, value);
}
return tree;
};
let tree = null;
data.forEach((item) => {
tree = insert(tree, item);
});
return tree;
};
const tree = generator(data);上面代码用generator 函数将 data 数组转成了一个排序二叉树 tree。将 tree 打印出来是这个样子：
以 9 节点为根节点的二叉树。打印的深度为叶子节点下一层，左右子树为空的话，就打印 null。
中序遍历找到第 k 个数
javascript
let index = 1;
const findKMin = (tree, k) => {
if (!tree || index > k) return null;
findKMin(tree.left, k, index + 1);
if (k == index) {
console.log("the k value is ", tree.value);
index++;
return;
}
index++;
findKMin(tree.right, k, index + 1);
};上面定义了一个函数findKMin，代码的逻辑是一个中序遍历递归算法的模版，在输出节点的值的部分，对 k 做了判断，如果 k!==index,那么就 index++，表示已经访问过了一个节点，接着继续访问下一个节点。如果k==index，说明我们找到了第 k 个数，那就将其打印出来。并且这里还要继续 index++；这里的目的是让 index 不再等于，防止访问到下一个节点的时候，还有k==index 的发生。
还有，在函数的开头做了index > k 的判断，如果index > k，就直接返回，说明已经找到了第 k 个数了。避免无谓的节点遍历
测试代码
javascript
findKMin(tree, 3);
// the k value is 4
findKMin(tree, 5);
// the k value is 7测试结果出来了，怎么验证这个结果的正确性呢？对比上面打印的树形结构吗，不要，这样效率太低了。可以直接将二叉树的中序遍历打印出来，然后数一数第 k 个数是不是上面输出的结果：
javascript
const printTree = (tree) => {
if (!tree) return null;
printTree(tree.left);
console.log(tree.value);
printTree(tree.right);
};
printTree(tree);
// 2
// 3
// 4
// 5
// 7
// 8
// 9
// 22
// 33
// 95上面就是二叉树的中序遍历输出序列了，可以看到，第 3 个数就是 4；第 5 个数就是 7。
没有问题
利用 count 变量
上面的解法没有用到题目中提到的 count 变量，算是另辟蹊径。如果用到 count，那就完全不同的解题思路了。
问题来了， count 可以为我们寻找第 k 个数提供什么样的帮助呢？
先思考一问题，对于二叉排序树，我想要找第 3 个数，而 root 节点的左子树（包含左子树的根节点，下同）有 6 个节点。那么，请问第 3 个树是在左子树中呢，还是右子树中呢？
肯定是左子树嘛，因为数数字，就是从左子树数过来的。
第二个问题，如果 root 的左子树只有两个节点，请问第 3 个数在哪？答案：第三个数就是 root 节点。还是因为中序遍历，中序遍历先左，再根，最后右子树。左子树数完了，下一个自然就是根节点了
第三个问题，如果 root 左子树只有一个节点，那么第 3 个数在哪？显而易见，在右子树中，并且是右子树的根节点，即 root.right;
现在相信你，一定可以理解 count 有什么用了。接下来就看看代码怎么写吧
准备数据
diff
const data = [9, 4, 5, 3, 2, 7, 8, 95, 33, 22];
const generator = (data) => {
const insert = (tree, value) => {
- if (!tree) return { value, left: null, right: null };
+ if (!tree) return { value, left: null, right: null, count: 0 };
if (value < tree.value) {
tree.left = insert(tree.left, value);
} else {
tree.right = insert(tree.right, value);
}
+ tree.count = (tree.left?.count || 0) + (tree.right?.count || 0);
+ if (tree.left) tree.count++;
+ if (tree.right) tree.count++;
return tree;
};
let tree = null;
data.forEach((item) => {
tree = insert(tree, item);
});
return tree;
};
const tree = generator(data);为了利用题目中提到的 count ，我们生成的排序二叉树中就需要有 count，所以我加了上面三行，用来统计每个节点 count。
对于统计 count 节点想了解更多，可以看这篇文章：🥳每日一练-统计二叉树节点的数量-JS简易版 - 掘金
方便理解，我更新打印的树形结构：
打印节点的同时，还将节点 count 打印处理
找到第 k 个数
javascript
const findMin2 = (tree, k) => {
if (k < 1 || !tree) return null;
if (tree.left) {
if (tree.left.count + 1 == k - 1) return tree;
if (tree.left.count + 1 < k - 1) return findMin2(tree.right, k - tree.left.count - 1);
if (tree.left.count + 1 > k - 1) return findMin2(tree.left, k);
}
k--;
if (k == 1) return tree;
return findMin2(tree.right, k);
};上面定义了一个findMin2，用来找到 tree 中的第 k 个数。代码很简单，逻辑自然不难，我们来捋一捋：
- 如果左子树存在，就看左子树的
count，如果tree.left.count+1 == k - 1，就说明当前访问的tree节点就是第k个数
为什么要加 1，因为左子树本身算一个
- 如果
tree.left.count +1 < k - 1，说明即使加上当前访问的tree节点，也还没有到第k个数。就需要在tree.right上继续找，递归调用findMin2，并且更新传入的 k。第二个参数是k - tree.left.count - 1，意思是前面已经数了tree.left.count + 1了，接下来只需要找第k - tree.left.count - 1数就可以了。
假设要找第 6 个数，已经数了 4 个，那么接下来只要再找 2 个数，就可以找到整棵树的第 6 个数了
- 如果
tree.left.count+1 > k - 1, 说明第k个数仍然在左子树，递归调用findMin2，并且不用更新传入的k。因为我还需要在左子树中找前k个 - 如果左子树不存在，就看看右子树的情况了。从上面可以看到，到了右子树，就有 k 减少的情况。这时候看看
k是否等于 1，如果等于 1，那么当前的tree就是要找的节点。为什么？可以将这个函数的入参k， 视作由上一个调用者findMin2调用时，传入k - tree.left.count - 1得来的 - 查看自身是否为第
k个数，需要消耗一个数，所以k-- - 如果
k!==1， 那就接着往右子树寻找，递归调用findMin2，并且更新传入的 k。第二个参数是k - 1，因为已经数过了当前的tree，所以只需要再找第k-1个就可以了 - 在函数的开头，有
k < 1的判断，和!tree一样，是对入参合法性的校验，并无他意。
测试代码
javascript
const node = findMin2(tree, 6);
console.log("node: ", node.value);
// node: 8
const node2 = findMin2(tree, 8);
console.log("node: ", node2.value);
// node: 33可以和上面的中序遍历输出对比，看看是否正确
总结
这篇文章分享了一道王道考研数据结构参考书的算法题，求出第 k 个数不难，感觉难的是用上题目中给出 count 这个条件。文中给出了清晰的解题思路，并且还对 count 做了详细生动的讲解，意在帮助大家理解。
你觉得这篇文章怎么样？我每天都会分享一篇算法小练习，喜欢就点赞+关注吧