题解 | #牛牛的二叉树问题#

/**
 * struct TreeNode {
 *	int val;
 *	struct TreeNode *left;
 *	struct TreeNode *right;
 *	TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 * };
 */
 #include <queue>
 #include <algorithm>

class mypair {
	int first;
	double second;
public:
	bool operator<(const mypair& other) const {
		return this->second < other.second;  // left > right
	}
    mypair(int a, double b):first(a), second(b){};
    double get_second() const {return second;};
    int get_first() const {return first;};
};

class Solution {
public:
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     * 
     * @param root TreeNode类 
     * @param target double浮点型 
     * @param m int整型 
     * @return int整型vector
     */
    void bfs(TreeNode* root, priority_queue<mypair>& heapmax, int& m, double& target){
        if (root == nullptr) return;
        if(heapmax.size() < m) {
            heapmax.emplace(root->val, abs(root->val - target));
        }
        else{
            // compair
            double temp = abs(root->val - target);
            if(temp < heapmax.top().get_second()){
                heapmax.pop();
                heapmax.emplace(root->val, temp);
            }
        }
        bfs(root->left, heapmax, m, target);
        bfs(root->right, heapmax, m, target);
     }
    vector<int> findClosestElements(TreeNode* root, double target, int m) {
        // write code here
        priority_queue<mypair> heapmax;
        vector<int> res;
        bfs(root, heapmax, m, target);
        for(int i = 0; i < m; ++i){
            res.push_back(heapmax.top().get_first());
            heapmax.pop();
        }
        sort(res.begin(), res.end());
        return res;
    }
};