【每日算法】用树建图，求距离目标节点为k的点集的两种方式

本文正在参加「Python主题月」，详情查看活动链接

题目描述

这是 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 / 迭代加深」找到距离为

k

的节点集。

而树是一类特殊的图，我们可以通过将二叉树转换为图的形式，再进行「BFS / 迭代加深」。

由于二叉树每个点最多有

2

个子节点，点和边的数量接近，属于稀疏图，因此我们可以使用「邻接表」的形式进行存储。

建图方式为：对于二叉树中相互连通的节点（root 与 root.left、root 和 root.right），建立一条无向边。

建图需要遍历整棵树，使用 DFS 或者 BFS 均可。

由于所有边的权重均为

1

，我们可以使用「BFS / 迭代加深」找到从目标节点 target 出发，与目标节点距离为

k

的节点，然后将其添加到答案中。

一些细节：利用每个节点具有唯一的值，我们可以直接使用节点值进行建图和搜索。

建图 + `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 进行建图的复杂度为 $O(n)$
空间复杂度： $O(n)$

建图 + 迭代加深

由「基本分析」，可写出「建图 + 迭代加深」的实现。

迭代加深的形式，我们只需要结合题意，搜索深度为

k

的这一层即可。

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 进行建图的复杂度为 $O(n)$
空间复杂度： $O(n)$

答疑

评论区不少小伙伴对 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 原题链接和其他优选题解。

【每日算法】用树建图，求距离目标节点为k的点集的两种方式

题目描述

基本分析

建图 + `BFS`

建图 + 迭代加深

答疑

最后

近期文章

近期评论

标签

热门

文章归档

分类目录

功能

题目描述

基本分析

建图 + BFS

建图 + 迭代加深

答疑

最后

近期文章

近期评论

标签

热门

文章归档

分类目录

功能

建图 + `BFS`