tags: Leetcode medium tree dfs bfs python c++

863. All Nodes Distance K in Binary Tree

[題目連結:] https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/description/

題目:

解題想法:

• 此題為給一target目標節點，返回所有跟target相距k步的節點

  • 各點值皆不重複
  • 可往上拜訪父點，也可往下拜訪子點

• DFS+BFS拜訪

  • 所需工具:
    • dic= collections.defaultdict(list)
      • {key:某點的值 ; value:與他相連的所有鄰居的值}
      • ex: dic[5]={3,6,2}
    • que=collections.deque()
      • bfs用來紀錄目前走到哪層
    • visited=set()
      • 紀錄已拜訪過的節點
  • 額外函式 dfs(parent, child, dic):
    分別將父點與子點關係紀錄於dic中，並遞迴紀錄
  ​​​​def dfs(self, parent, child, dic):
  ​​​​    if parent and child:
  ​​​​        dic[parent.val].append(child.val)
  ​​​​        dic[child.val].append(parent.val)
  ​​​​    if child.left:
  ​​​​        self.dfs(child, child.left, dic)
  ​​​​    if child.right:
  ​​​​        self.dfs(child, child.right, dic)
  ​​​​    return
  

  • 主函式: 目標擴展k層
    • 第一層:5
    • 第二層:3,2,6
    • 第三層(res):1 7 4
  ​​​​for _ in range(k): #擴展K層
  ​​​​    n=len(que)
  ​​​​    for i in range(len(que)):
  ​​​​        curNode=que.popleft()
  ​​​​        for neighber in dic[curNode]:
  ​​​​            if neighber not in visited:
  ​​​​                visited.add(neighber)
  ​​​​                que.append(neighber)
  ​​​​return list(que)
  

Python:

from collections import defaultdict, deque
# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    def distanceK(self, root, target, k):
        """
        :type root: TreeNode
        :type target: TreeNode
        :type k: int
        :rtype: List[int]
        """
        dic=defaultdict(list)
        que=deque()
        self.dfs(None, root, dic)
        que.append(target.val)
        visited=set({target.val}) #紀錄已拜訪過的點
        #擴展K層    第一層:5    第二層:3,2,6  第三層(res):1 7 4
        for _ in range(k):
            n=len(que)
            for i in range(len(que)):
                curNode=que.popleft()
                for neighber in dic[curNode]:
                    if neighber not in visited:
                        visited.add(neighber)
                        que.append(neighber)
        return list(que)
        
    #key:某點的值 value:與他相連的所有鄰居的值 ex: dic[5]={3,6,2}
    def dfs(self, parent, child, dic):
        if parent and child:
            dic[parent.val].append(child.val)
            dic[child.val].append(parent.val)
        if child.left:
            self.dfs(child, child.left, dic)
        if child.right:
            self.dfs(child, child.right, dic)
        return

C++:

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        unordered_map<int, vector<int>> dic;
        queue<int> que;
        set<int> visited;
        dfs(nullptr, root, dic);
        que.push(target->val);
        visited.insert(target->val);
        //run k times
        for (int i=0; i<k; i++){
            int n=que.size();
            for (int j=0; j<n; j++){
                int curNode=que.front(); que.pop();
                for (auto neighber: dic[curNode]){
                     //not visited yet
                    if (!visited.count(neighber)){
                        visited.insert(neighber);
                        que.push(neighber);
                    }
                }
            }
        }
        vector<int> res;
        int num=que.size();
        for (int i=0; i<num; i++){
            int tmp=que.front(); que.pop();
            res.push_back(tmp);
        }
        return res;
    }

    void dfs(TreeNode* parent, TreeNode* child, unordered_map<int, vector<int>>& dic){
        if (parent && child){
            dic[parent->val].push_back(child->val);
            dic[child->val].push_back(parent->val);
        }
        if (child->left) dfs(child, child->left, dic);
        if (child->right) dfs(child, child->right, dic);
        return;
    }
};