[LeetCode] 863. All Nodes Distance K in Binary Tree 二叉樹距離爲K的全部結點

We are given a binary tree (with root node root), a target node, and an integer value K.

Return a list of the values of all nodes that have a distance K from the target node.  The answer can be returned in any order.

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2

Output: [7,4,1]

Explanation:
The nodes that are a distance 2 from the target node (with value 5)
have values 7, 4, and 1.

Note that the inputs "root" and "target" are actually TreeNodes.
The descriptions of the inputs above are just serializations of these objects.

Note:

1. The given tree is non-empty.
2. Each node in the tree has unique values 0 <= node.val <= 500.
3. The target node is a node in the tree.
4. 0 <= K <= 1000.

這道題給了咱們一棵二叉樹，一個目標結點 target，還有一個整數K，讓返回全部跟目標結點 target 相距K的結點。咱們知道在子樹中尋找距離爲K的結點很容易，由於只須要一層一層的向下遍歷便可，難點就在於符合題意的結點有多是祖先結點，或者是在旁邊的兄弟子樹中，這就比較麻煩了，由於二叉樹只有從父結點到子結點的路徑，反過來就不行。既然沒有，咱們就手動建立這樣的反向鏈接便可，這樣樹的遍歷問題就轉爲了圖的遍歷（其實樹也是一種特殊的圖）。創建反向鏈接就是用一個 HashMap 來來創建每一個結點和其父結點之間的映射，使用先序遍歷創建好全部的反向鏈接，而後再開始查找和目標結點距離K的全部結點，這裏須要一個 HashSet 來記錄全部已經訪問過了的結點。

在遞歸函數中，首先判斷當前結點是否已經訪問過，是的話直接返回，不然就加入到 visited 中。再判斷此時K是否爲0，是的話說明當前結點已是距離目標結點爲K的點了，將其加入結果 res 中，而後直接返回。不然分別對當前結點的左右子結點調用遞歸函數，注意此時帶入 K-1，這兩步是對子樹進行查找。以前說了，還得對父結點，以及兄弟子樹進行查找，這是就體現出創建的反向鏈接 HashMap 的做用了，若當前結點的父結點存在，咱們也要對其父結點調用遞歸函數，並一樣帶入 K-1，這樣就能正確的找到全部知足題意的點了，參見代碼以下：

解法一：

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        if (!root) return {};
        vector<int> res;
        unordered_map<TreeNode*, TreeNode*> parent;
        unordered_set<TreeNode*> visited;
        findParent(root, parent);
        helper(target, K, parent, visited, res);
        return res;
    }
    void findParent(TreeNode* node, unordered_map<TreeNode*, TreeNode*>& parent) {
        if (!node) return;
        if (node->left) parent[node->left] = node;
        if (node->right) parent[node->right] = node;
        findParent(node->left, parent);
        findParent(node->right, parent);
    }
    void helper(TreeNode* node, int K, unordered_map<TreeNode*, TreeNode*>& parent, unordered_set<TreeNode*>& visited, vector<int>& res) {
        if (visited.count(node)) return;
        visited.insert(node);
        if (K == 0) {res.push_back(node->val); return;}
        if (node->left) helper(node->left, K - 1, parent, visited, res);
        if (node->right) helper(node->right, K - 1, parent, visited, res);
        if (parent[node]) helper(parent[node], K - 1, parent, visited, res);
    }
};

既然是圖的遍歷，那就也可使用 BFS 來作，爲了方便起見，咱們直接創建一個鄰接鏈表，即每一個結點最多有三個跟其相連的結點，左右子結點和父結點，使用一個 HashMap 來創建每一個結點和其相鄰的結點數組之間的映射，這樣就幾乎徹底將其看成圖來對待了，創建好鄰接鏈表以後，原來的樹的結構都不須要用了。既然是 BFS 進行層序遍歷，就要使用隊列 queue，還要一個 HashSet 來記錄訪問過的結點。在 while 循環中，若K爲0了，說明當前這層的結點都是符合題意的，就把當前隊列中全部的結點加入結果 res，並返回便可。不然就進行層序遍歷，取出當前層的每一個結點，並在鄰接鏈表中找到和其相鄰的結點，若沒有訪問過，就加入 visited 和 queue 中便可。記得每層遍歷完成以後，K要自減1，參見代碼以下：

解法二：

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        if (!root) return {};
        vector<int> res;
        unordered_map<TreeNode*, vector<TreeNode*>> m;
        queue<TreeNode*> q{{target}};
        unordered_set<TreeNode*> visited{{target}};
        findParent(root, NULL, m);
        while (!q.empty()) {
            if (K == 0) {
                for (int i = q.size(); i > 0; --i) {
                    res.push_back(q.front()->val); q.pop();
                }
                return res;
            }
            for (int i = q.size(); i > 0; --i) {
                TreeNode *t = q.front(); q.pop();
                for (TreeNode *node : m[t]) {
                    if (visited.count(node)) continue;
                    visited.insert(node);
                    q.push(node);
                }
            }
            --K;
        }
        return res;
    }
    void findParent(TreeNode* node, TreeNode* pre, unordered_map<TreeNode*, vector<TreeNode*>>& m) {
        if (!node) return;
        if (m.count(node)) return;
        if (pre) {
            m[node].push_back(pre);
            m[pre].push_back(node);
        }
        findParent(node->left, node, m);
        findParent(node->right, node, m);
    }
};

其實這道題也能夠不用 HashMap，不創建鄰接鏈表，直接在遞歸中完成全部的需求，真正體現了遞歸的博大精深。在進行遞歸以前，咱們要先判斷一個 corner case，那就是當 K==0 時，此時要返回的就是目標結點值自己，能夠直接返回。不然就要進行遞歸了。這裏的遞歸函數跟以前的有所不一樣，是須要返回值的，這個返回值表示的含義比較複雜，若爲0，表示當前結點爲空或者當前結點就是距離目標結點爲K的點，此時返回值爲0，是爲了進行剪枝，使得不用對其左右子結點再次進行遞歸。當目標結點正好是當前結點的時候，遞歸函數返回值爲1，其餘的返回值爲當前結點離目標結點的距離加1。還須要一個參數 dist，其含義爲離目標結點的距離，注意和遞歸的返回值區別，這裏不用加1，且其爲0時候不是爲了剪枝，而是真不知道離目標結點的距離。

在遞歸函數中，首先判斷若當前結點爲空，則直接返回0。而後判斷 dist 是否爲k，是的話，說目標結點距離當前結點的距離爲K，是符合題意的，須要加入結果 res 中，並返回0，注意這裏返回0是爲了剪枝。不然判斷，若當前結點正好就是目標結點，或者已經遍歷過了目標結點（表現爲 dist 大於0），那麼對左右子樹分別調用遞歸函數，並將返回值分別存入 left 和 right 兩個變量中。注意此時應帶入 dist+1，由於是先序遍歷，若目標結點以前被遍歷到了，那麼說明目標結點確定不在當前結點的子樹中，當前要往子樹遍歷的話，確定離目標結點又遠了一些，須要加1。若當前結點不是目標結點，也還沒見到目標結點時，一樣也須要對左右子結點調用遞歸函數，但此時 dist 不加1，由於不肯定目標結點的位置。若 left 或者 right 值等於K，則說明目標結點在子樹中，且距離當前結點爲K（爲啥呢？由於目標結點自己是返回1，因此當左右子結點返回K時，和當前結點距離是K）。接下來判斷，若當前結點是目標結點，直接返回1，這個前面解釋過了。而後再看 left 和 right 的值是否大於0，若 left 值大於0，說明目標結點在左子樹中，咱們此時就要對右子結點再調用一次遞歸，而且 dist 帶入 left+1，同理，若 right 值大於0，說明目標結點在右子樹中，咱們此時就要對左子結點再調用一次遞歸，而且 dist 帶入 right+1。這兩步很重要，是之因此能不創建鄰接鏈表的關鍵所在。若 left 大於0，則返回 left+1，若 right 大於0，則返回 right+1，不然就返回0，參見代碼以下：

解法三：

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        if (K == 0) return {target->val};
        vector<int> res;
        helper(root, target, K, 0, res);
        return res;
    }
    int helper(TreeNode* node, TreeNode* target, int k, int dist, vector<int>& res) {
        if (!node) return 0;
        if (dist == k) {res.push_back(node->val); return 0;}
        int left = 0, right = 0;
        if (node->val == target->val || dist > 0) {
            left = helper(node->left, target, k, dist + 1, res);
            right = helper(node->right, target, k, dist + 1, res);
        } else {
            left = helper(node->left, target, k, dist, res);
            right = helper(node->right, target, k, dist, res);
        }
        if (left == k || right == k) {res.push_back(node->val); return 0;}
        if (node->val == target->val) return 1;
        if (left > 0) helper(node->right, target, k, left + 1, res);
        if (right > 0) helper(node->left, target, k, right + 1, res);
        if (left > 0 || right > 0) return left > 0 ? left + 1 : right + 1;
        return 0;
    }
};

Github 同步地址:

https://github.com/grandyang/leetcode/issues/863

參考資料：

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/discuss/143752/JAVA-Graph-%2B-BFS

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/discuss/143775/very-easy-to-understand-c%2B%2B-solution.

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/discuss/143886/Java-O(1)-space-excluding-recursive-stack-space

LeetCode All in One 題目講解彙總(持續更新中...)