863. All Nodes Distance K in Binary Tree

link

Method 1

find the node and the path then start from path node which is less than k, if the node is in pass, then skip

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:

        self.arr = []

        def find(node, current):
            if self.arr:
                return
            if not node:
                return
            if node == target:
                current.append(node)
                self.arr = current.copy()
                return

            current.append(node)

            find(node.left, current)
            find(node.right, current)

            current.pop()
            

        find(root, [])
        

        n = len(self.arr)

        self.arr2 = []

        def dfs(node, c, nxt):
            if not node:
                return
            # print(node.val, c)
            if node == nxt:
                return
            if c == 0:
                
                self.arr2.append(node.val)
                return
            
            dfs(node.left, c-1, nxt)
            dfs(node.right, c-1, nxt)
            

        for i in range(n-1, max(-1, n-k-2), -1):
            node = None
            if i != n - 1:
                node = self.arr[i+1]
            dfs(self.arr[i],  k - (n-1 - i), node)




        return self.arr2

Method2

Remember the parent, can also add a dict and save the info(like a connected graph)

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        def dfs(node, parent):
            if not node:
                return
            
            node.parent = parent
            dfs(node.left, node)
            dfs(node.right, node)
            
        dfs(root, None)
        queue = deque([target])
        seen = {target}
        distance = 0
        
        while queue and distance < k:
            current_length = len(queue)
            for _ in range(current_length):
                node = queue.popleft()
                for neighbor in [node.left, node.right, node.parent]:
                    if neighbor and neighbor not in seen:
                        seen.add(neighbor)
                        queue.append(neighbor)
            
            distance += 1
        
        return [node.val for node in queue]