题解 | #距离是K的二叉树节点#

题目里没说明的一点是输出结果的顺序, 这个还得自己看样例来猜。

整体思路是先遍历一遍二叉树, 一边记录下每个节点的父节点, 一边找到目标节点。

找到目标节点之后, 先左后右最后根的顺序开始递归找距离为k的节点即可。

import java.util.*;

public class Solution {
    Map<TreeNode, TreeNode> parentMap = new HashMap<>();

    ArrayList<Integer> res = new ArrayList<>();

    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     * @param root TreeNode类
     * @param target int整型
     * @param k int整型
     * @return int整型一维数组
     */
    public ArrayList<Integer> distanceKnodes(TreeNode root, int target, int k) {
        ArrayDeque<TreeNode> queue = new ArrayDeque<>();
        TreeNode targetNode = null;
        queue.addFirst(root);
        while (!queue.isEmpty()) {
            TreeNode curRoot = queue.pollLast();
            TreeNode cLeft = curRoot.left;
            TreeNode cRight = curRoot.right;
            if (cLeft != null) {
                if (cLeft.val == target) {
                    targetNode = cLeft;
                }
                queue.addFirst(cLeft);
                parentMap.put(cLeft, curRoot);
            }
            if (cRight != null) {
                if (cRight.val == target) {
                    targetNode = cRight;
                }
                queue.addFirst(cRight);
                parentMap.put(cRight, curRoot);
            }
        }
        if (root.val == target) { // 如果根就是目标节点的话,上边是没遍历到的..
            targetNode = root;
        }
        if (targetNode == null) {
            return new ArrayList<>();
        }
        HashSet<TreeNode> usedNodesSet = new HashSet<>();
        findK(targetNode, usedNodesSet, k);
        return res;
    }

    private void findK(TreeNode curNode, HashSet<TreeNode> usedNodesSet, int k) {
        if (curNode == null || usedNodesSet.contains(curNode)) {
            return;
        }
        usedNodesSet.add(curNode);
        if (k == 0) {
            res.add(curNode.val);
            return;
        }
        findK(curNode.left, usedNodesSet, k - 1);
        findK(curNode.right, usedNodesSet, k - 1);
        findK(parentMap.get(curNode), usedNodesSet, k - 1);
    }
}