863. All Nodes Distance K in Binary Tree (Medium) - TengnanYao/daily_leetcode GitHub Wiki

# 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]
        """
        def dfs(node, parent = None):
            if node:
                node.parent = parent
                dfs(node.left, node)
                dfs(node.right, node)
        dfs(root)
        queue = [target]
        seen = {target}
        distance = 0
        while queue:
            if distance == K:
                return [node.val for node in queue]
            newQueue = []
            for node in queue:
                seen.add(node)
                if node.left and node.left not in seen:
                    newQueue.append(node.left)
                if node.right and node.right not in seen:
                    newQueue.append(node.right)
                if node.parent and node.parent not in seen:
                    newQueue.append(node.parent)
            distance += 1
            queue = newQueue
        return []