Computer >> Hướng Dẫn Máy Tính >  >> Lập Trình >> C++

Tìm K giá trị gần nhất trong BST với giá trị mục tiêu trong C++

Cho BST và giá trị target (float), tìm k giá trị trong cây gần target nhất.

Ví dụ

BST: [4,2,5,1,3], target=3.714286, k=2 → Output: [4, 3]

Thuật toán: Hai stack (Inorder + Reverse Inorder)

  • Stack smaller: các nút < target, top là lớn nhất < target.
  • Stack larger: các nút ≥ target, top là nhỏ nhất ≥ target.
  • Lặp k lần: so sánh độ chênh lệch top hai stack, lấy gần hơn, sau đó đẩy nút con phù hợp.

Triển khai C++

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

struct TreeNode {
    int val;
    TreeNode *left, *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

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

int main() {
    // Build tree: 4,2,5,1,3
    TreeNode* root = new TreeNode(4);
    root->left = new TreeNode(2);
    root->right = new TreeNode(5);
    root->left->left = new TreeNode(1);
    root->left->right = new TreeNode(3);
    
    Solution ob;
    auto res = ob.closestKValues(root, 3.7142, 2);
    for (int x : res) cout << x << " "; // 4 3
    return 0;
}

Kết quả

4 3

Độ phức tạp: O(log n + k) trung bình, O(n) worst-case.