[leetcode] 863. All Nodes Distance K in Binary Tree @ python

原题

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

解法

先使用DFS, 构建conn字典将数的元素连接起来, 由于题目中既要向下找子树, 又要向上找, 所以我们要双向连接. 运行connect函数后, conn变成{3: [5,1], 5: [3, 6, 2], 2: [5, 7, 4], 1:[3, 0, 8]…}.

然后使用BFS, 构建当前bfs列表, 和已见过的seen集合, 然后我们做K次bfs搜索循环来找到距离K的所有节点, 搜索后将见过的节点与set取并集. 因此在第K次搜索时, bfs里的节点就是seen中没有的, 即与target节点距离K的节点了.

Time: O(n)
Space: O(n)

代码

# 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]
        """
        # dfs
        conn = collections.defaultdict(list)
        def connect(parent, child):
            if parent and child:
                conn[parent.val].append(child.val)
                conn[child.val].append(parent.val)                
            if child.left:
                connect(child, child.left)
            if child.right:
                connect(child, child.right)
        connect(None, root)
        # bfs
        bfs = [target.val]
        seen = set(bfs)
        for i in range(K):
            bfs = [y for x in bfs for y in conn[x] if y not in seen]
            seen = seen | set(bfs)
        return bfs