動機

不要被BST誤導阿!!! 一直往O(log(n))的方向去想

Problem

Given the root of a binary search tree, and an integer k, return the kth (1-indexed) smallest element in the tree.

 

Example 1:

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

Example 2:

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

 

Constraints:

  • The number of nodes in the tree is n.
  • 1 <= k <= n <= 104
  • 0 <= Node.val <= 104

 

Follow up: If the BST is modified often (i.e., we can do insert and delete operations) and you need to find the kth smallest frequently, how would you optimize?

Sol

做中序去找第k個

心得

  1. 0 是 False
  2. 沒有log(n),沒有log(n),沒有log(n)
def dfs(r,k,acc):
    if not r:
        return [None, acc]
    elif not r.left and not r.right:
        return [r.val,acc+1]
    else:
        ret = dfs(r.left,k,acc)
        # 0 is False in Python3 !!!!
        if ret[0] != None and k == ret[1]:
            return ret
        ret[0] = r.val
        ret[1] += 1
        if ret[0] != None and k == ret[1]:
            return ret
        return dfs(r.right,k,ret[1])
            
class Solution:
    def kthSmallest(self, r: TreeNode, k: int) -> int:
        return dfs(r,k,0)[0]