Giá trị BST gần nhất II trong C++

Cho cây tìm kiếm nhị phân (BST) và giá trị target (số thực), tìm k giá trị trong BST gần target nhất. k luôn hợp lệ, k ≤ tổng số nút.

Ví dụ:

target = 3.714286, k = 2 → [4, 3]

Phương pháp giải

Sử dụng 2 stack: smaller (giá trị < target) và larger (giá trị ≥ target). pushSmaller: đi sang phải nếu nút < target, ngược lại sang trái. pushLarger: ngược lại. Lặp k lần: lấy từ stack nào gần target hơn, thêm vào kết quả, đẩy con thích hợp của nút vừa lấy vào stack tương ứng.

Ví dụ chương trình

#include <bits/stdc++.h>
using namespace std;

void print_vector(vector v) {
    cout << "[";
    for (int i = 0; i < v.size(); i++) cout << v[i] << ", ";
    cout << "]" << endl;
}

class TreeNode {
public:
    int val;
    TreeNode *left, *right;
    TreeNode(int data) : val(data), left(NULL), right(NULL) {}
};

void insert(TreeNode** root, int val) {
    queue q;
    q.push(*root);
    while (q.size()) {
        TreeNode* temp = q.front(); q.pop();
        if (!temp->left) {
            if (val != NULL) temp->left = new TreeNode(val);
            else temp->left = new TreeNode(0);
            return;
        } else q.push(temp->left);
        if (!temp->right) {
            if (val != NULL) temp->right = new TreeNode(val);
            else temp->right = new TreeNode(0);
            return;
        } else q.push(temp->right);
    }
}

TreeNode* make_tree(vector v) {
    TreeNode* root = new TreeNode(v[0]);
    for (int i = 1; i < v.size(); i++) insert(&root, v[i]);
    return root;
}

class Solution {
public:
    vector closestKValues(TreeNode* root, double target, int k) {
        vector ret;
        stack smaller, larger;
        pushLarger(root, larger, target);
        pushSmaller(root, smaller, target);
        
        while (k--) {
            if (!smaller.empty() && (larger.empty() || 
                abs(target - smaller.top()->val) < abs(target - larger.top()->val))) {
                TreeNode* curr = smaller.top(); smaller.pop();
                ret.push_back(curr->val);
                pushSmaller(curr->left, smaller, target);
            } else {
                TreeNode* curr = larger.top(); larger.pop();
                ret.push_back(curr->val);
                pushLarger(curr->right, larger, target);
            }
        }
        return ret;
    }
    
    void pushSmaller(TreeNode* node, stack& st, double target) {
        while (node) {
            if (node->val < target) { st.push(node); node = node->right; }
            else node = node->left;
        }
    }
    
    void pushLarger(TreeNode* node, stack& st, double target) {
        while (node) {
            if (node->val >= target) { st.push(node); node = node->left; }
            else node = node->right;
        }
    }
};

int main() {
    Solution ob;
    vector v = {4,2,5,1,3};
    TreeNode* root = make_tree(v);
    print_vector(ob.closestKValues(root, 3.7142, 2));
    return 0;
}

Đầu vào

BST: [4,2,5,1,3], target=3.7142, k=2

Đầu ra

[4, 3]