Fork me on GitHub

Closest Binary Search Tree Value II

Posted on

Description

https://leetcode.com/problems/closest-binary-search-tree-value-ii/description/

Solution

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
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
struct cmp {
    bool operator () (pair<float, int>& A, pair<float, int>& B) {
        return A.first < B.first;
    }
};

class Solution {
public:
    vector<int> closestKValues(TreeNode* root, double target, int k) {
        
        priority_queue<pair<float, int>, vector<pair<float, int>>, cmp> pq;
        inOrderTraverse(root, target, k, pq);
        vector<int> ret;
        while (!pq.empty()) {
            ret.push_back(pq.top().second);
            pq.pop();
        }
        return ret;
    }
    
    void inOrderTraverse(TreeNode* root, double target, int k, 
         priority_queue<pair<float, int>, vector<pair<float, int>>, cmp>& pq) {
        if (root == NULL) return;
        inOrderTraverse(root->left, target, k, pq);
        auto pa = make_pair(abs(root->val - target), root->val);
        if (pq.size() >= k) {
            if (pa.first < pq.top().first) {
                pq.pop();
                pq.push(pa);
            } 
        }else {
            pq.push(pa);
        }
        inOrderTraverse(root->right, target, k, pq);
    }
};