Problem description:

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.
Example:

Input: root = [4,2,5,1,3], target = 3.714286, and k = 2

1
2
3
4
5
    4
   / \
  2   5
 / \
1   3

Output: [4,3]
Follow up:
Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?

Solution:

the smaller stack will pop the next available largest element that is smaller then target, and larger stack will pop the next available smallest element that is larger or equal to the target

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
# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def closestKValues(self, root: TreeNode, target: float, k: int) -> List[int]:
        res = []
        pre, suc = deque([]), deque([])
        
        while root:
            if root.val > target:
                suc.append(root)
                root = root.left
            else:
                pre.append(root)
                root = root.right
        
        def getPredecessor(stack):
            top = stack.pop()
            p = top.left
            while p:
                stack.append(p)
                p = p.right
            return top
        def getSucceesor(stack):
            top = stack.pop()
            p = top.right
            while p:
                stack.append(p)
                p = p.left
            return top
        
        
        while k:
            k -= 1
            if pre and not suc:
                p = getPredecessor(pre)
                res.append(p.val)
            elif not pre and suc:
                p = getSucceesor(suc)
                res.append(p.val)
            elif pre and suc and abs(pre[-1].val-target) <= abs(suc[-1].val-target):
                p = getPredecessor(pre)
                res.append(p.val)
            elif pre and suc and abs(pre[-1].val-target) >= abs(suc[-1].val-target):
                p = getSucceesor(suc)
                res.append(p.val)
        return res

time complexity: $O(logn)$
space complexity: $O(logn)$

Use dfs+ heap to get overall minimum distance nodes.

Notice python is min heap, every time when heap size exceed k , would pop out the smallest element, so we add negative to keep the closest node(it would be largest distance after negative).

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
# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def closestKValues(self, root: TreeNode, target: float, k: int) -> List[int]:
        def dfs(root):
            if not root:
                return
            
            if len(heap) == k:
                heapq.heappushpop(heap, (-abs(root.val-target), root.val))
            else:
                heapq.heappush(heap, (-abs(root.val-target), root.val))
                
            dfs(root.left)
            dfs(root.right)
            
        heap, res = [], []
        dfs(root)
        for _ in range(k):
            res.append(heapq.heappop(heap)[1])
        return res

My first idea is to use priority queue to solve this question. Since every time we push an element into the priority_queue, it will automatically sort the elements inside the queue in an ascending order, we can track the size of priority_queue to determine whether if we need to pop out some elements.

The reason he uses max heap is that when pq.size() > k, we do pq.pop(). This pops out the largest value, which is equivalent to saying keeping the k smallest elements measured in terms of distance to the target value.

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
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<int> closestKValues(TreeNode* root, double target, int k) {
        priority_queue<pair<double, int>> pq;
        helper(root, target, k, pq);
        vector<int> res;
        while(!pq.empty()){
            res.push_back(pq.top().second);
            pq.pop();
        }
        return res;
    }
    
    void helper(TreeNode* root, double target, int k, priority_queue<pair<double, int>>& pq){
        if(!root) return;
        
        pq.push(make_pair(fabs(double(root->val)- target), root->val));
        
        if(pq.size()> k)
            pq.pop();
        
        helper(root->left, target, k, pq);
        helper(root->right, target, k, pq);
    }
};

time complexity: $O(nlogk)$
space complexity: $O(k)$
reference:
http://www.cnblogs.com/grandyang/p/5247398.html