863. All Nodes Distance K in Binary Tree - Medium

前往題目

想法

• 直接把target當root，執行DFS

思路

就差了一點，忘了從target開始得想辦法往parent那邊找

1. 把所有child的parent都紀錄起來
2. dfs尋找，只要再多加上找target的parents
3. 利用visited(hashset)防止重複加入node

visited是為了找parent的時候不重複加入已經檢查過的node

Code

討論區答案

class Solution {
    List<Integer> res = new ArrayList<>();
    Map<TreeNode, TreeNode> childParentMap = new HashMap<>();
    Set<TreeNode> visited = new HashSet<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        // Build the children to parent map
        findParent(root);
        dfs(target, k);
        return res;
    }

    private void findParent(TreeNode node) {
        if (node == null) {
            return;
        }
        if (node.left != null) {
            childParentMap.put(node.left, node);
            findParent(node.left);
        }
        if (node.right != null) {
            childParentMap.put(node.right, node);
            findParent(node.right);
        }
    }

    private void dfs(TreeNode node, int k) {
        // Base case
        if (node == null || visited.contains(node))
            return;
        // Mark as visited
        visited.add(node);
        if (k == 0) {
            res.add(node.val);
            return;
        }

        // Find through its children
        dfs(node.left, k - 1);
        dfs(node.right, k - 1);
        // Find through its parents
        dfs(childParentMap.get(node), k - 1);

        return;
    }
}