题目描述
这是 LeetCode 上的 863. 二叉树中所有距离为 K 的结点 ，难度为 中等。
Tag : 「图论 BFS」、「图论 DFS」、「二叉树」
给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 K 。
返回到目标结点 target 距离为 K 的所有结点的值的列表。 答案可以以任何顺序返回。
示例 1：
输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2输出：[7,4,1]
解释：
所求结点为与目标结点（值为 5）距离为 2 的结点，
值分别为 7，4，以及 1
注意，输入的 "root" 和 "target" 实际上是树上的结点。 上面的输入仅仅是对这些对象进行了序列化描述。
提示：
- 给定的树是非空的。
- 树上的每个结点都具有唯一的值 0 <= node.val <= 500 。
- 目标结点 target 是树上的结点。
- 0 <= K <= 1000.
基本分析
显然，如果题目是以图的形式给出的话，我们可以很容易通过「BFS / 迭代加深」找到距离为 的节点集。
而树是一类特殊的图，我们可以通过将二叉树转换为图的形式，再进行「BFS / 迭代加深」。
由于二叉树每个点最多有 个子节点，点和边的数量接近，属于稀疏图，因此我们可以使用「邻接表」的形式进行存储。
建图方式为：对于二叉树中相互连通的节点（root 与 root.left、root 和 root.right），建立一条无向边。
建图需要遍历整棵树，使用 DFS 或者 BFS 均可。
由于所有边的权重均为 ，我们可以使用 「BFS / 迭代加深」 找到从目标节点 target 出发，与目标节点距离为 的节点，然后将其添加到答案中。
一些细节：利用每个节点具有唯一的值，我们可以直接使用节点值进行建图和搜索。
建图 + BFS
由「基本分析」，可写出「建图 + BFS」的实现。
Java 代码：
class Solution {int N = 1010, M = N * 2;
int[] he = new int[N], e = new int[M], ne = new int[M];
int idx;
void add(int a, int b) {
e[idx] = b;
ne[idx] = he[a];
he[a] = idx++;
}
boolean[] vis = new boolean[N];
public List<Integer> distanceK(TreeNode root, TreeNode t, int k) {
List<Integer> ans = new ArrayList<>();
Arrays.fill(he, -1);
dfs(root);
Deque<Integer> d = new ArrayDeque<>();
d.addLast(t.val);
vis[t.val] = true;
while (!d.isEmpty() && k >= 0) {
int size = d.size();
while (size-- > 0) {
int poll = d.pollFirst();
if (k == 0) {
ans.add(poll);
continue;
}
for (int i = he[poll]; i != -1 ; i = ne[i]) {
int j = e[i];
if (!vis[j]) {
d.addLast(j);
vis[j] = true;
}
}
}
k--;
}
return ans;
}
void dfs(TreeNode root) {
if (root == null) return;
if (root.left != null) {
add(root.val, root.left.val);
add(root.left.val, root.val);
dfs(root.left);
}
if (root.right != null) {
add(root.val, root.right.val);
add(root.right.val, root.val);
dfs(root.right);
}
}
}
Python 3 代码：
class Solution:# 根据数据范围最多有 501 个点，每个点最多有 2 条无向边（两个子节点）
N = 510
M = N * 4
def distanceK(self, root: TreeNode, t: TreeNode, k: int) -> List[int]:
he = [-1] * self.N
e = [0] * self.M
ne = [0] * self.M
idx = 0
vis = [False] * self.N
def add(a, b):
nonlocal idx
e[idx] = b
ne[idx] = he[a]
he[a] = idx
idx += 1
def dfs(root):
if not root:
return
if root.left:
add(root.val, root.left.val)
add(root.left.val, root.val)
dfs(root.left)
if root.right:
add(root.val, root.right.val)
add(root.right.val, root.val)
dfs(root.right)
ans = []
dfs(root)
d = deque([t.val])
vis[t.val] = True
while d and k >= 0:
size = len(d)
while size > 0:
poll = d.popleft()
size -= 1
if not k:
ans.append(poll)
continue
i = he[poll]
while i != -1:
j = e[i]
if not vis[j]:
d.append(j)
vis[j] = True
i = ne[i]
k -= 1
return ans
- 时间复杂度：通过DFS 进行建图的复杂度为；通过BFS 找到距离为的节点，复杂度为。整体复杂度为
- 空间复杂度：
建图 + 迭代加深
由「基本分析」，可写出「建图 + 迭代加深」的实现。
迭代加深的形式，我们只需要结合题意，搜索深度为 的这一层即可。
Java 代码：
class Solution {int N = 1010, M = N * 2;
int[] he = new int[N], e = new int[M], ne = new int[M];
int idx;
void add(int a, int b) {
e[idx] = b;
ne[idx] = he[a];
he[a] = idx++;
}
boolean[] vis = new boolean[N];
public List<Integer> distanceK(TreeNode root, TreeNode t, int k) {
List<Integer> ans = new ArrayList<>();
Arrays.fill(he, -1);
dfs(root);
vis[t.val] = true;
find(t.val, k, 0, ans);
return ans;
}
void find(int root, int max, int cur, List<Integer> ans) {
if (cur == max) {
ans.add(root);
return ;
}
for (int i = he[root]; i != -1; i = ne[i]) {
int j = e[i];
if (!vis[j]) {
vis[j] = true;
find(j, max, cur + 1, ans);
}
}
}
void dfs(TreeNode root) {
if (root == null) return;
if (root.left != null) {
add(root.val, root.left.val);
add(root.left.val, root.val);
dfs(root.left);
}
if (root.right != null) {
add(root.val, root.right.val);
add(root.right.val, root.val);
dfs(root.right);
}
}
}
Python 3 代码：
class Solution:# 根据数据范围最多有 501 个点，每个点最多有 2 条无向边（两个子节点）
N = 510
M = N * 4
def distanceK(self, root: TreeNode, t: TreeNode, k: int) -> List[int]:
he = [-1] * self.N
e = [0] * self.M
ne = [0] * self.M
idx = 0
vis = [False] * self.N
def add(a, b):
nonlocal idx
e[idx] = b
ne[idx] = he[a]
he[a] = idx
idx += 1
def dfs(root):
if not root:
return
if root.left:
add(root.val, root.left.val)
add(root.left.val, root.val)
dfs(root.left)
if root.right:
add(root.val, root.right.val)
add(root.right.val, root.val)
dfs(root.right)
def find(root, m, cur):
if cur == m:
ans.append(root)
return
i = he[root]
while i != -1:
j = e[i]
if not vis[j]:
vis[j] = True
find(j, m, cur + 1)
i = ne[i]
ans = []
dfs(root)
vis[t.val] =
- 时间复杂度：通过DFS 进行建图的复杂度为；通过迭代加深找到距离为的节点，复杂度为。整体复杂度为
- 空间复杂度：
答疑
评论区不少小伙伴对 add 的存图方式有疑问，这里集中回答一下 ~
这是一种在图论中十分常见的存图方式，可直接当作模板进行背过，与数组存储单链表的实现一致。
首先 idx 是用来对边进行编号的，然后对存图用到的几个数组作简单解释：
- he 数组：存储是某个节点所对应的边的集合（链表）的头结点；
- e 数组：由于访问某一条边指向的节点；
- ne 数组：由于是以链表的形式进行存边，该数组就是用于找到下一条边。
因此当我们想要遍历所有由 a 点发出的边时，可以使用如下方式：
for (int i = he[a]; i != -1; i = ne[i]) {int j = e[i]; // 存在由 a 指向 j 的边
}
另外，在 LeetCode 评论区 @Meteordream 小姐姐给出了很好的解释：
数组 he 的下标表示结点，值是一个索引 ind，e[ind] 表示 对应一条边，ne[ind] 表示下一个连接结点的索引，假设与 结点a 相连的结点有 b, c, 那么通过 he[a]取得一个索引 ind1 后，通过 e[ind1] = b 可以得到与 a 相连的第一个结点是 b，然后通过 ne[ind1] 可以获得下一个结点的索引 ind2 ，通过 e[ind2] = c 可以得到与 a 相连的第二个结点是 c，最后 ne[ind2] = -1 说明没有下一个结点了
add函数采用链表的头插法，假设 结点a 已经有一个相连的结点 b，那么就有 he[a]=ind, e[ind]=b ，此时再给 a 增加一个相连的结点 c，那么就要建立由b的索引到新结点c的索引 ne[new_ind] = he[a] = ind ，然后新建一条边 e[new_ind], 最后更新 he[a] = new_ind ，就完成了由 a -> b 到 a -> c -> b 的添加操作
可以理解为 he 是邻接表的表头，key是结点val是一个指向存有相邻结点的链表头指针，e是链表结点的val即相邻结点，ne是链表结点的next指针
如果还有疑问的小伙伴，可以带着「链式前向星存图」关键字进行搜索学习哦 ~
最后
这是我们「刷穿 LeetCode」系列文章的第 No.863 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先将所有不带锁的题目刷完。
在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。
为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：github.com/SharingSour… 。
在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。