863. All Nodes Distance K in Bin

2019-05-19  本文已影响0人  尚无花名

We are given a binary tree (with root node root), a target node, and an integer value K.
Return a list of the values of all nodes that have a distance K from the target node. The answer can be returned in any order.

Example 1:
Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2
Output: [7,4,1]
Explanation:
The nodes that are a distance 2 from the target node (with value 5)
have values 7, 4, and 1.
这道题不能一个recursion写完。
要分两个recursion来写。
先把这个点找出来. 找的时候同时记录从parent to 这个node的路径上的所有点, 可以存在一个stack里面。
这个可以用一个recursion来做。
然后从stack里面一个一个取,从每个点开始找相应的点。
这里面你会发现最后一个点和其他的点不一样。
最后一个点你可以既找它的左孩子,又可以找它的右孩子。但其他点只能找左孩子右孩子中间的其中一个。
所以我这里面用了一个prev来记录上一个node,如果上一个node是当前是node的左孩子,我就只从当前node右子树去找。

class Solution {
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        List<Integer> ans = new ArrayList<>();
        Deque<TreeNode> deque = new ArrayDeque<>();
        if (root == null || target == null) return ans;
        
        findPath(root, deque, target);
        
        TreeNode prev = null;
        int dist = 0;
        
        while (!deque.isEmpty()) {
            TreeNode node = deque.pop();
            if (dist == K) {
                ans.add(node.val);
                return ans;
            } 
            if (node.left != prev) findAllNode(node.left, K - dist - 1, ans);
            if (node.right != prev) findAllNode(node.right, K - dist - 1, ans);
            prev = node;
            dist++;
        }
        return ans;
    }
    public void findAllNode(TreeNode node, int level, List<Integer> ans) {
        if (node == null || level < 0) return;
        if (level == 0) { 
            ans.add(node.val);
            return;
        }
        findAllNode(node.left, level - 1, ans);
        findAllNode(node.right, level - 1, ans);
    }
    
    private boolean findPath(TreeNode node, Deque<TreeNode> deque, TreeNode target) {
        if (node == null) return false;
        deque.push(node);
        if (node == target) return true;
        if (findPath(node.left, deque, target)) return true;
        if (findPath(node.right, deque, target)) return true;
        deque.pop();
        return false;
    }
}
上一篇下一篇

猜你喜欢

热点阅读