230. Kth Smallest Element in a BST

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

Example 1:

Input: root = [3,1,4,null,2], k = 1
3
/ \
1 4
\
  2
Output: 1

Example 2:

Input: root = [5,3,6,2,4,null,null,1], k = 3
5
/ \
3 6
/ \
2 4
/
1
Output: 3

Follow up: What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?

Constraints:

The number of elements of the BST is between 1 to 10^4. You may assume k is always valid, 1 ≤ k ≤ BST's total elements.

Analyze

思路: 中序遍历维护计数, 当计数值到达 k 时, 当前节点的 val 即为题目要求的第 k 小的值元素值。

/**
* Definition for a binary tree node.
* function TreeNode(val, left, right) {
* this.val = (val===undefined ? 0 : val)
* this.left = (left===undefined ? null : left)
* this.right = (right===undefined ? null : right)
* }
*/
/**
* @param {TreeNode} root
* @param {number} k
* @return {number}
*/
var kthSmallest = function(root, k) {
const countResult = { value: 0 }
return analyzeCount(root, k, countResult)
}
var analyzeCount = function(node, k, count) {
if (!node) return null
const pickLeft = analyzeCount(node.left, k, count)
if (typeof pickLeft === 'number') return pickLeft
count.value = count.value + 1
if (k === count.value) {
return node.val
}
const pickRight = analyzeCount(node.right, k, count)
if (typeof pickRight === 'number') return pickRight
}