算法练习 kth smallest element in BST with known subtree size,在已知子树大小的二叉搜索树里找第k小的元素

Sunday, Mar 29, 2020 | 2 minute read | Updated at Sunday, Mar 29, 2020

最简单暴力的就是直接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))