863. All Nodes Distance K in Binary Tree

1/**
2 * Definition for a binary tree node.
3 * class TreeNode {
4 *     val: number
5 *     left: TreeNode | null
6 *     right: TreeNode | null
7 *     constructor(val?: number, left?: TreeNode | null, right?: TreeNode | null) {
8 *         this.val = (val===undefined ? 0 : val)
9 *         this.left = (left===undefined ? null : left)
10 *         this.right = (right===undefined ? null : right)
11 *     }
12 * }
13 */
14
15function distanceK(root: TreeNode | null, target: TreeNode | null, k: number): number[] {
16    const parent: Map<TreeNode, TreeNode | null> = new Map();
17    function dfs(node: TreeNode | null, par: TreeNode | null): void {
18        if (!node) return;
19        parent.set(node, par);
20        dfs(node.left, node);
21        dfs(node.right, node);
22    }
23    dfs(root, null);
24
25    const res: number[] = [];
26    const vis: Set<TreeNode> = new Set();
27
28    function dfs2(node: TreeNode | null, from: TreeNode | null, dist: number) {
29        if (!node || vis.has(node)) return;
30        vis.add(node);
31        if (dist === 0) {
32            res.push(node.val);
33            return;
34        }
35
36        dfs2(node.left, node, dist - 1);
37        dfs2(node.right, node, dist - 1);
38        dfs2(parent.get(node) || null, node, dist - 1);
39    }
40
41    dfs2(target, null, k);
42    return res;
43};