LeetCode 230. Kth Smallest Element in a BST

Posted on | In |
Words count in article: 255 | Reading time ≈ 1

题目

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

Example 1:

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

Example 2:

1
2
3
4
5
6
7
8
9
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.

思路

Inorder traversal of BST is an array sorted in the ascending order.

Follow up:
we could add a variable to the TreeNode to record the size of the left subtree. When insert or delete a node in the left subtree, we increase or decrease it by 1. So we could know whether the kth smallest element is in the left subtree or in the right subtree by compare the size with k.

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution(object):
    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        return self.inOrder(root)[k-1]
    
    
    def inOrder(self, root):
        return self.inOrder(root.left) + [root.val] + self.inOrder(root.right) if root else []