Leetcode 863. All Nodes Distance K in Binary Tree (python)

题目

解析：

这道题看到固定距离寻找这种限制条件，第一个想到的就应该是BFS。所以这道题目变成了两个部分：

• 将树转化为图，每个节点指向parent节点和左右子节点
• 从target节点出发，BFS寻找特定距离的点
• 值得注意的是，这边给的target是节点，所以不需要额外去寻找target节点
  这边将树转化为图有两种做法，一种是利用树的性质，将parent直接作为树节点的属性，代码如下：

class Solution(object):
    def distanceK(self, root, target, K):
        """
        :type root: TreeNode
        :type target: TreeNode
        :type K: int
        :rtype: List[int]
        """
        def add_par(node,par):
            if not node:
                return
            node.par = par
            add_par(node.left,node)
            add_par(node.right,node)
        
        add_par(root,None)
        
        ans = []
        q = collections.deque()
        q.append((target,0))
        visited = set()
        visited.add(target.val)
        while q:
            curr,dis = q.popleft()
            if dis == K:
                ans.append(curr.val)
            for nei in (curr.left,curr.right,curr.par):
                if nei and nei.val not in visited:
                    q.append((nei,dis+1))
                    visited.add(nei.val)
        return ans 

而第二种才是更加标准的做法。将树转化为图通常的表达方式，也就是字典

class Solution(object):
    def distanceK(self, root, target, K):
        """
        :type root: TreeNode
        :type target: TreeNode
        :type K: int
        :rtype: List[int]
        """
        def build_graph(node,par):
            if not node:
                return
            if node.left:
                graph[node].append(node.left)
            if node.right:
                graph[node].append(node.right)
            if par:
                graph[node].append(par)
            build_graph(node.left,node)
            build_graph(node.right,node)
        graph = collections.defaultdict(list)
        build_graph(root,None)
        ans = []
        q = collections.deque()
        q.append((target,0))
        visited = set()
        visited.add(target.val)
        while q:
            node,dis = q.popleft()
            if dis == K:
                ans.append(node.val)
            for nei in graph[node]:
                if nei.val not in visited:
                    q.append((nei,dis+1))
                    visited.add(nei.val)
        return ans

二刷解法：

在bfs里面每次pop出相同距离的所有node，这样超过目标距离的节点就无需访问

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        # transform tree into graph
        g = collections.defaultdict(list)
        def dfs(node):
            if node.left:
                g[node.val].append(node.left.val)
                g[node.left.val].append(node.val)
                dfs(node.left)
            if node.right:
                g[node.val].append(node.right.val)
                g[node.right.val].append(node.val)
                dfs(node.right)
        
        dfs(root)
        # print(g)
        # bfs to find the nodes at k distance
        q = collections.deque()
        visited = {}
        q.append(target.val)
        visited[target.val] = True
        dist = 0
        while q:
            if dist == k:
                return list(q)
            # pop out all the nodes at the current distance
            for _ in range(len(q)):
                curr = q.popleft()
                for neigh in g[curr]:
                    if neigh not in visited:
                        q.append(neigh)
                        visited[neigh] = True
            dist += 1
        return []

c++版本

class Solution {
public:
    void helper(TreeNode* node,unordered_map<int,vector<int>>& g){
        if(!node) return;
        if(node->left){
            g[node->val].push_back(node->left->val);
            g[node->left->val].push_back(node->val);
            helper(node->left,g);
        }
        if(node->right){
            g[node->val].push_back(node->right->val);
            g[node->right->val].push_back(node->val);
            helper(node->right,g);
        }
        return;
    }
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        unordered_map<int,vector<int>> g;
        helper(root,g);
        queue<int> q;
        unordered_set<int> visited;
        q.push(target->val);
        visited.insert(target->val);
        int dis = 0;
        vector<int> ans;
        while(!q.empty()){
            int n = q.size();
            for(int i=0;i<n;i++){
                int curr = q.front();
                // cout << curr << " " << dis << endl;
                q.pop();
                if(dis == k){
                    ans.push_back(curr);
                }
                for(auto& neigh : g[curr]){
                    if(visited.count(neigh)) continue;
                    q.push(neigh);
                    visited.insert(neigh);
                }
            }
            if(dis == k){
                return ans;
            }
            dis++;
        }
        return ans;
    }
};

时间复杂度：O(n)
空间复杂度：O(n)
n为节点个数