863. All Nodes Distance K in Binary Tree

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/

二叉树中找与给定结点相距K的所有结点。二叉树是一种特殊的图,因此可以把二叉树存成无向图然后用BFS找所有相距K的点。还有可以利用二叉树的性质做DFS,对target的每个ancestor(与t距离为d)收集它在另一个子树中与ancestor距离为K-d-1的所有结点。

/*
 * @lc app=leetcode id=863 lang=cpp
 *
 * [863] All Nodes Distance K in Binary Tree
 */

// @lc code=start
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<int> res;
    void collect(TreeNode *cur, int d) {
        if (cur == nullptr) return;
        if (d == 0) {
            res.push_back(cur->val);
            return;
        } 
        collect(cur->left, d - 1);
        collect(cur->right, d - 1);
    }

    int dist(TreeNode *cur, TreeNode *target, const int K) {
        if (cur == nullptr) return -1;
        if (cur == target) {
            collect(cur, K);
            return 0;
        }
        int l = dist(cur->left, target, K);
        if (l != -1) {
            ++l;
            if (K - l > 0) {
                collect(cur->right, K - l - 1);
            } else if (K == l) {
                res.push_back(cur->val);
            }
            return l;
        }
        int r = dist(cur->right, target, K);
        if (r != -1) {
            ++r;
            if (K - r > 0) {
                collect(cur->left, K - r - 1);
            } else if (K == r) {
                res.push_back(cur->val);
            }
            return r;
        }
        return -1;
    }

    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        dist(root, target, K);
        return res;
    }
};
// @lc code=end

Last updated