863. All Nodes Distance K in Binary Tree

发表于

每日一题 2019 - 02 - 24

题目：

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:

1
2
3
4
5
6
7
8
9
10
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.

Note that the inputs "root" and "target" are actually TreeNodes.
The descriptions of the inputs above are just serializations of these objects.

1

Note:

  1. The given tree is non-empty.
  2. Each node in the tree has unique values 0 <= node.val <= 500.
  3. The target node is a node in the tree.
  4. 0 <= K <= 1000.

解法：

参考网上大佬的做法，如下：

  • 我们需要做若干次深度优先遍历。
  • 第一次需要定位 target 的具体位置，并按顺序记录其到根结点的路径。路径中的每个点，用一个二元组表示，分别为点的指针，和一个布尔值。其中布尔值为 true 代表 target 是在当前结点的左侧，否则在右侧。路径不包括 target 本身。
  • 然后，我们顺着路径，从底向上。如果路径上当前的点已经和 target 距离为 K，则保存当前结点，然后退出循环。否则，如果布尔值为 true，则从右侧开始深度优先遍历找距离 target 为 K 的结点，否则从左侧开始找。
  • 最后，从 target 开始深度优先遍历找答案。

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
class Solution {
public:
    void find(TreeNode *rt, int cur, int K, vector<int>& ans) {
        if (rt == nullptr || cur > K)
            return;
        if (cur == K) {
            ans.push_back(rt -> val);
            return;
        }
        find(rt -> left, cur + 1, K, ans);
        find(rt -> right, cur + 1, K, ans);
    }

    bool findTarget(TreeNode *rt, TreeNode *target, vector<pair<TreeNode*, bool>>& path) {
        if (rt == nullptr)
            return false;
        if (rt == target)
            return true;

        if (findTarget(rt -> left, target, path)) {
            path.emplace_back(make_pair(rt, true));
            return true;
        }
        if (findTarget(rt -> right, target, path)) {
            path.emplace_back(make_pair(rt, false));
            return true;
        }
        return false;
    }

    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        vector<pair<TreeNode*, bool>> path;
        findTarget(root, target, path);
        vector<int> ans;

        for (int i = 0; i < path.size(); i++) {
            if (i + 1 == K) {
                ans.push_back(path[i].first -> val);
                break;
            }
            if (path[i].second == true)
                find(path[i].first -> right, i + 2, K, ans);
            else
                find(path[i].first -> left, i + 2, K, ans);
        }

        find(target, 0, K, ans);
        return ans;
    }
};
0%
undefined