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

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



Leave a Comment

Your email address will not be published.