動機

轉成graph錯了嗎

Problem

Given the root of a binary tree, the value of a target node target, and an integer k, return an array of the values of all nodes that have a distance k from the target node.

You can return the answer in any order.

 

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2Output: [7,4,1]Explanation: The nodes that are a distance 2 from the target node (with value 5) have values 7, 4, and 1.

Example 2:

Input: root = [1], target = 1, k = 3Output: []

 

Constraints:

  • The number of nodes in the tree is in the range [1, 500].
  • 0 <= Node.val <= 500
  • All the values Node.val are unique.
  • target is the value of one of the nodes in the tree.
  • 0 <= k <= 1000

Sol

轉成graph，就可以跑dfs了

而solution上的做法是幫node加parent pointer，這樣就可以往上走了

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        gh = defaultdict(set)
        
        def buildGh(r):
            if r:
                dirs = [r.left, r.right]
                [gh[r.val].add(x.val) or gh[x.val].add(r.val) for x in dirs if x]
                [buildGh(x) for x in dirs if x]
        
        seen = set()
        ret = []
        def dfs(r,k):
            if r not in seen:
                seen.add(r)
                if k == 0:
                    ret.append(r)
                else:
                    for x in gh[r]:
                        dfs(x,k-1)
                #seen.remove(r)
        
        buildGh(root)
        dfs(target.val,k)
        return ret