1. 程式人生 > >LeetCode-All Nodes Distance K in Binary Tree

LeetCode-All Nodes Distance K in Binary Tree

一、Description

題目描述:給定一個二叉樹,一個target結點,一個K,找出所有離target結點的距離為K的所有結點,返回一個list。

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2

Output: [7,4,1]


二、Analyzation

  1. 從target結點開始dfs(target, k),找出所有target子結點中所有滿足題意的結點,加入到list中。
  2. 找出target的父結點,執行dfs(target, k - 1),找出所有沒有訪問過的滿足題意的結點。
  3. 找出target的父結點的父結點和父結點的其它子結點的父結點,執行dfs(target, k - 2)...以此內推,當k = 0時將當前訪問結點加入到list中,遞歸回溯。

三、Accepted code

class Solution {
    Map<TreeNode, TreeNode> map = new HashMap<>();
    Set<TreeNode> set = new HashSet<>();
    List<Integer> list = new ArrayList<>();
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        if (root == null) {
            return list;
        }
        findParent(root);
        dfs(target, K);
        return list;
    }
    public void findParent(TreeNode root) {
        if (root.left != null) {
            map.put(root.left, root);
            findParent(root.left);
        }
        if (root.right != null) {
            map.put(root.right, root);
            findParent(root.right);
        }
    }
    public void dfs(TreeNode node, int k) {
        if (set.contains(node)) {
            return ;
        }
        set.add(node);
        if (k == 0) {
            list.add(node.val);
        }
        if (node.left != null) {
            dfs(node.left, k - 1);
        }
        if (node.right != null) {
            dfs(node.right, k - 1);
        }
        TreeNode parent = map.get(node);
        if (parent != null) {
            dfs(parent, k - 1);
        }
    }
}