最簡單暴力的就是直接dfs做in order traversal,記錄現在是第幾個元素,當等於k時就返回。這樣時間複雜度是O(N),空間O(logN)。但是這樣就沒有利用到子樹大小已知的屬性。
最優解法是利用子樹大小信息推出當前節點的index,如果比k大,那麼目標在當前節點左邊,如果比k小則在當前節點右邊,如果等於k,就找到了。這樣可以把時間減到O(logN),空間則是O(1),優於暴力解法不少。
推出當前節點的index,可以用一個變數記錄range,也就是子樹可能的最小和最大index。
一道挺有意思的題目。
Python
# 50 (10)
# / \
# 30 (4) 70 (5)
# / \ / \
# 20 (2) 40 (1) 60 (3) 80 (1)
# / / \
# 5 (1) 55 (1) 65 (1)
class TreeNode:
def __init__(self, val):
self.size = 1
self.val = val
self.left = None
self.right = None
def build_sample_tree():
node50 = TreeNode(50)
node30 = TreeNode(30)
node20 = TreeNode(20)
node40 = TreeNode(40)
node70 = TreeNode(70)
node60 = TreeNode(60)
node80 = TreeNode(80)
node5 = TreeNode(5)
node55 = TreeNode(55)
node65 = TreeNode(65)
node50.left = node30
node50.right = node70
node50.size = 10
node30.left = node20
node30.right = node40
node30.size = 4
node20.left = node5
node20.size = 2
node70.left = node60
node70.right = node80
node70.size = 5
node60.left = node55
node60.right = node65
node60.size = 3
return node50
def dfs(root, range, n):
if n < range[0] or n > range[1]:
return None
if root is None:
return None
l_size = 0
if root.left is not None:
l_size = root.left.size
idx = range[0] + l_size
if idx == n:
return root.val
elif idx < n:
range[0] = idx + 1
return dfs(root.right, range, n)
else:
range[1] = idx - 1
return dfs(root.left, range, n)
def nsmallest(root, n):
return dfs(root, [1, float('inf')], n)
root = build_sample_tree()
for i in range(12):
print '{}-th smallest: {}'.format(i, nsmallest(root, i))
