https://leetcode.com/problems/kth-smallest-element-in-a-bst/
Kth Smallest Element in a BST
Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.
Note:
You may assume k is always valid, 1 ≤ k ≤ BST’s total elements.
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?
Hint:
- Try to utilize the property of a BST.
- What if you could modify the BST node’s structure?
- The optimal runtime complexity is O(height of BST).
Code
# Definition for a binary tree node.
# class TreeNode(object):
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
class Solution(object):
def kthSmallest(self, root, k):
"""
:type root: TreeNode
:type k: int
:rtype: int
"""
for v in self.inorder(root):
if k == 1:
return v
else:
k -= 1
def inorder(self, root):
if root:
for v in self.inorder(root.left):
yield v
yield root.val
for v in self.inorder(root.right):
yield v
Idea
We use depth first in order traverse of BST to read element from the smallest to the largest. We can count when each element is read, until we read k element. Here we use `yield` instead of `return` in the inorder function to be more memory thrifty.
Locating the smallest element will take O(log N) time, plus you read O(k) element before you return.
Reference
https://leetcode.com/discuss/44731/pythonic-approach-with-generator