-
Notifications
You must be signed in to change notification settings - Fork 0
/
kth smallest emelment in BST.py
60 lines (55 loc) · 1.61 KB
/
kth smallest emelment in BST.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
'''
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).
'''
class Solution(object):
def kthSmallest(self, root, k):
"""
:type root: TreeNode
:type k: int
:rtype: int
"""
current = root
stack = []
while current:
while current:
stack.append(current)
current = current.left
current = stack.pop()
k -= 1
if k == 0:
return current.val
while not current.right:
current = stack.pop()
k -= 1
if k == 0:
return current.val
current = current.right
'''
Simplify
'''
class Solution(object):
def kthSmallest(self, root, k):
"""
:type root: TreeNode
:type k: int
:rtype: int
"""
current = root
stack = []
while current or stack:
while current:
stack.append(current)
current = current.left
current = stack.pop()
k -= 1
if k == 0:
return current.val
current = current.right