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
}