[LeetCode]All Nodes Distance K in Binary Tree

这道题有几种做法，首先如果我们知道每个节点对应的父节点的话，我们直接从target节点开始bfds即可。那么我们可以先dfs一次建立每个节点和对应的父节点的关系，然后我们从target节点开始bfs找到左右距离为k的节点即可。假设输入的树的size为n，那么时间复杂度为O(n)，空间复杂度O(n)。
另一种方法，所有和target距离为k的节点可以分为几类：

• 在target为根的子树中并且和target距离为k
• 是target其中一个ancestor，和target距离恰好为k
• 在target的ancestor中的另一个子树中。比如target是在ancestor的左子树，并且target和ancestor的距离为i(i < k)。那么ancestor的右子树中就可能存在一个和ancestor距离为(k - i)的节点，其和target的距离正好为i + k - i = k

所以我们可以先遍历一遍树，寻找所有可能存在和target距离为k的节点存在的子树的根，之后在子树中寻找和根距离为i的所有节点即可。时间复杂度O(n)，空间复杂度O(n)，代码如下：

	class Solution {
	public:
	vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
	vector<int> res;
	vector<pair<TreeNode*, int>> cands;
	cands.emplace_back(target, K);
	preprocess(root, target, cands, K, res);
	for (auto& cand : cands)
	collectInSubtree(cand.first, cand.second, res);
	return res;
	}
	private:
	int preprocess(TreeNode* root, TreeNode* target, vector<pair<TreeNode*, int>>& cands, int K, vector<int>& res)
	{
	if (!root)return -1;
	if (root == target)return 1;
	int leftD = preprocess(root->left, target, cands, K, res);
	int rightD = preprocess(root->right, target, cands, K, res);
	if (leftD == -1 && rightD == -1)
	return -1;
	else if (leftD == -1)
	{
	if (rightD == K)
	res.push_back(root->val);
	else if (rightD < K)
	if(root->left)cands.emplace_back(root->left, K - rightD - 1);
	return rightD + 1;
	}
	else
	{
	if (leftD == K)
	res.push_back(root->val);
	else if (leftD < K)
	if(root->right)cands.emplace_back(root->right, K - leftD - 1);
	return leftD + 1;
	}
	}
	void collectInSubtree(TreeNode* root, int dist, vector<int>& res)
	{
	if (!root)return;
	if (!dist)
	{
	res.push_back(root->val);
	return;
	}
	collectInSubtree(root->left, dist - 1, res);
	collectInSubtree(root->right, dist - 1, res);
	}
	};

view raw All Nodes Distance K in Binary Tree.cpp hosted with ❤ by GitHub