Closest Binary Search Tree Value II

题目描述

Given a non-empty binary search tree and a target value, find k values in the BST that are closest to the target.

Note: Given target value is a floating point. You may assume k is always valid, that is: k ≤ total nodes. You are guaranteed to have only one unique set of k values in the BST that are closest to the target.

Follow up:

Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?

解题方法

在遍历这个tree的过程中，要不断地更新距离target的距离，时刻替换。 而一般的data structure都没有比较的作用，只有heap可以维护一个有比较的 structure, 那么用来比较的value就是跟target的距离。

follow u

利用Inorder得到pre-decessors的性质来做，逆inorder就是得到successors

  • 找到这个target的closest node
  • 对于这个node找到pre-decessors和successors
  • 然后用类似merge sort的方法找到k个最近的值

Solution

class Solution(object):
    def closestKValues(self, root, target, k):
        """
        :type root: TreeNode
        :type target: float
        :type k: int
        :rtype: List[int]
        """

        heap = []
        stack = []

        stack.append(root)
        while stack:
            cur = stack.pop()
            if cur.right:
                stack.append(cur.right)
            if cur.left:
                stack.append(cur.left)

            if len(heap) < k:
                heapq.heappush(heap, (abs(cur.val - target) * -1, cur.val))
            elif (abs(cur.val - target) * -1) > heap[0][0]:
                heapq.heappop(heap)
                heapq.heappush(heap, (abs(cur.val - target) * -1, cur.val))

        return [node[1] for node in heap]

Reference