題目網址：https://leetcode.cn/problems/all-nodes-distance-k-in-binary-tree/

題意：給一 BT 的 root, 和一目標點 target, 以及一整數 k。

返回所有距離目標點 target 為 k 的 node, 可以用任何順序返回。

tree 中每個 node.val 都不同。

Solution：

想法：利用 DFS + BFS

  • 先用 DFS 建立無向圖
  • 再以 target 為中心向外做 BFS

e.g. target = 5, k = 2

  • 以 5 為中心, 首先取出 3, 6, 2
  • 然後 3 會取出 5, 1
  • 依此類推…

這邊必須小心, 因為我們建立的是無向圖, child 的 adj list 中是會有 parent 的
以上面的為例, 3 向外擴散取出 5 跟 1。其中 1 是答案, 但 5 不是
因此, 得用一個 visited 來記錄拜訪過的 node

  • 若已經在 visited 裡則跳過
  • 否則, 將其加入 visited 中

這樣一來, 5 最一開始就會被加入到 visited 中
就算後面的 node 有跟 5 相鄰, 也不會把 5 加入到 res 中

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        buildGraph(root);

        vector<int> res;
        unordered_set<TreeNode*> visited;
        queue<TreeNode*> q;
        q.emplace(target);

        while (!q.empty()) {
            const int len = q.size();

            // 取出當前 level
            for (int i = 0; i < len; ++i) {
                const auto cur = q.front();
                q.pop();

                // 若已經拜訪過, 則跳過
                if (visited.find(cur) != visited.end()) {
                    continue;
                }

                // 否則, 將 cur 設為已拜訪
                visited.emplace(cur);

                // 為距離 target 的第 k 層
                if (k == 0) {
                    res.emplace_back(cur->val);
                } else { // 否則, 將其 adjacent node 加入到下一層
                    for (const auto& v : adj[cur]) {
                        q.emplace(v);
                    }
                }
            }

            --k; // 做完一層, 層數減一
        }

        return res;
    }

private:
    unordered_map<TreeNode*, vector<TreeNode*>> adj; // adjacency list

    void buildGraph(TreeNode* root){
        if (!root) {
            return;
        }

        if (root->left) {
            adj[root].emplace_back(root->left);
            adj[root->left].emplace_back(root);
            buildGraph(root->left);
        }

        if (root->right) {
            adj[root].emplace_back(root->right);
            adj[root->right].emplace_back(root);
            buildGraph(root->right);
        }
    }
};
  • time：$O(n)$ ➔ 每個 node 皆被遍歷常數次
  • space：$O(n)$ ➔ 撇除要返回的 res, 在 while loop 迭代的過程中, q 的元素個數不超過 n