【每日算法】用樹建圖，求距離目標節點爲 k 的點集的兩種方式：「建圖 + BFS」&「建圖 + 迭代加深」 ｜Python 主題月

本文正在參加「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)$；經過 BFS 找到距離 $target$ 爲 $k$ 的節點，複雜度爲 $O(n)$。總體複雜度爲 $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)$；經過迭代加深找到距離 $target$ 爲 $k$ 的節點，複雜度爲 $O(n)$。總體複雜度爲 $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 原題連接和其餘優選題解。