147 - 230 二叉树搜索树中第k小的元素

题目

给定一个二叉搜索树，编写一个函数 kthSmallest 来查找其中第 k 个最小的元素。

说明： 你可以假设 k 总是有效的，1 ≤ k ≤ 二叉搜索树元素个数。

示例 1:

输入: root = [3,1,4,null,2], k = 1 3 / 1 4 2 输出: 1

示例 2:

输入: root = [5,3,6,2,4,null,null,1], k = 3 5 / 3 6 / 2 4 / 1 输出: 3

进阶： 如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化 kthSmallest 函数？

解答

想到的方法一，是变成数组，然后用小顶堆或者快速搜索。

方法二就是中序遍历，数k个。

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        stack = []
        p = root
        count = 0
        while p is not None or stack:
            while p is not None:
                stack.append(p)
                p = p.left
            p = stack.pop()
            count += 1
            if count == k:
                return p.val
            p = p.right

Runtime: 48 ms, faster than 95.99% of Python3 online submissions for Kth Smallest Element in a BST.

Memory Usage: 16.6 MB, less than 98.18% of Python3 online submissions for Kth Smallest Element in a BST.

var kthSmallest = function(root, k) {
  const stack = []
  let p = root
  let count = 0
  while (p || stack) {
    while (p) {
      stack.push(p)
      p = p.left
    }
    p = stack.pop()
    count++
    if (count === k) {
      return p.val
    }
    p = p.right
  }
};

Runtime: 56 ms, faster than 98.94% of JavaScript online submissions for Kth Smallest Element in a BST.

Memory Usage: 39.4 MB, less than 63.64% of JavaScript online submissions for Kth Smallest Element in a BST.

func kthSmallest(root *TreeNode, k int) int {
    stack := []*TreeNode{}
    count := 0
    p := root
    for p != nil || len(stack) != 0 {
        for p != nil {
            stack = append(stack, p)
            p = p.Left
        }
        p = stack[len(stack)-1]
        stack = stack[:len(stack)-1]
        count++
        fmt.Println("count", count)
        if count == k {
            return p.Val
        }
        p = p.Right
    }
    return -1
}

Runtime: 12 ms, faster than 70.69% of Go online submissions for Kth Smallest Element in a BST.

Memory Usage: 5.8 MB, less than 10.00% of Go online submissions for Kth Smallest Element in a BST.