Tất cả nút ở khoảng cách K trong cây nhị phân trong C++

Giới thiệu

Cho cây nhị phân, nút đích target và số nguyên K. Tìm tất cả các nút có khoảng cách bằng K từ target. Khoảng cách là số cạnh trên đường đi.

Thuật toán: DFS + BFS

1. DFS từ gốc để xây dựng map parent[node] = cha.
2. BFS từ target với trạng thái (node, level):
   • Nếu level == K → thêm vào kết quả.
   • Duyệt 3 hướng: trái, phải, cha (nếu chưa visited).

Cài đặt C++

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

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

class Solution {
    unordered_map<TreeNode*, TreeNode*> parent;
    void dfs(TreeNode* node, TreeNode* par = nullptr) {
        if (!node) return;
        parent[node] = par;
        dfs(node->left, node);
        dfs(node->right, node);
    }
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        dfs(root);
        queue<pair<TreeNode*, int>> q;
        unordered_set<TreeNode*> visited;
        q.push({target, 0});
        visited.insert(target);
        vector<int> ans;
        while (!q.empty()) {
            auto [node, lvl] = q.front(); q.pop();
            if (lvl == k) ans.push_back(node->val);
            if (lvl == k) continue;
            if (node->left && !visited.count(node->left)) {
                q.push({node->left, lvl+1});
                visited.insert(node->left);
            }
            if (node->right && !visited.count(node->right)) {
                q.push({node->right, lvl+1});
                visited.insert(node->right);
            }
            if (parent[node] && !visited.count(parent[node])) {
                q.push({parent[node], lvl+1});
                visited.insert(parent[node]);
            }
        }
        return ans;
    }
};

// Hàm hỗ trợ tạo cây từ mảng (level order)
TreeNode* buildTree(vector<int> v) { /* ... */ }

int main() {
    Solution ob;
    vector<int> v = {3,5,1,6,2,0,8,-1,-1,7,4}; // -1 đại diện cho null
    TreeNode* root = buildTree(v);
    TreeNode* target = root->left; // node 5
    auto res = ob.distanceK(root, target, 2);
    for (int x : res) cout << x << ' '; // Output: 7 4 1
    return 0;
}

Phân tích

• Thời gian: O(N) – DFS và BFS mỗi nút tối đa 1 lần.
• Không gian: O(N) – map parent, queue, visited set.

Đây là giải pháp chuẩn cho LeetCode 863 "All Nodes Distance K in Binary Tree".