93 lines
1.9 KiB
C++
93 lines
1.9 KiB
C++
#include <unordered_map>
|
|
#include <vector>
|
|
#include <deque>
|
|
|
|
using namespace std;
|
|
|
|
/*
|
|
二叉树节点定义
|
|
*/
|
|
struct TreeNode {
|
|
int val;
|
|
TreeNode* left;
|
|
TreeNode* right;
|
|
};
|
|
|
|
/*
|
|
图的邻接表表示
|
|
key : 树节点指针
|
|
value : 与该节点相邻的节点列表
|
|
*/
|
|
unordered_map<TreeNode*, vector<TreeNode*>> adj;
|
|
|
|
/*
|
|
BFS 访问标记
|
|
防止在“树转无向图”后出现来回走
|
|
*/
|
|
unordered_map<TreeNode*, bool> visited;
|
|
|
|
/*
|
|
将二叉树转换为无向图
|
|
本质：父 <-> 子 建立双向边
|
|
*/
|
|
void TreeNodeToGraph(TreeNode* root) {
|
|
if (root == nullptr) return;
|
|
|
|
// 如果有左孩子，建立 root <-> left 的双向边
|
|
if (root->left != nullptr) {
|
|
adj[root].push_back(root->left);
|
|
adj[root->left].push_back(root);
|
|
|
|
TreeNodeToGraph(root->left);
|
|
}
|
|
|
|
// 如果有右孩子，建立 root <-> right 的双向边
|
|
if (root->right != nullptr) {
|
|
adj[root].push_back(root->right);
|
|
adj[root->right].push_back(root);
|
|
|
|
TreeNodeToGraph(root->right);
|
|
}
|
|
}
|
|
|
|
/*
|
|
从 target 出发做 BFS
|
|
返回：与 target 距离恰好为 step 的所有节点
|
|
*/
|
|
deque<TreeNode*> BFS(TreeNode* target, int step) {
|
|
deque<TreeNode*> q;
|
|
q.push_back(target);
|
|
|
|
visited[target] = true;
|
|
|
|
int dist = 0;
|
|
|
|
while (!q.empty()) {
|
|
int cnt = q.size(); // 当前层节点数
|
|
|
|
// 如果已经到达目标层，直接返回这一层的所有节点
|
|
if (dist == step) {
|
|
return q;
|
|
}
|
|
|
|
// 扩展当前层
|
|
for (int i = 0; i < cnt; i++) {
|
|
TreeNode* node = q.front();
|
|
q.pop_front();
|
|
|
|
// 遍历邻接节点
|
|
for (TreeNode* nxt : adj[node]) {
|
|
if (!visited[nxt]) {
|
|
visited[nxt] = true;
|
|
q.push_back(nxt);
|
|
}
|
|
}
|
|
}
|
|
|
|
dist++;
|
|
}
|
|
|
|
// 如果 step 超过树高度，返回空
|
|
return {};
|
|
}
|