3067. 在带权树网络中统计可连接服务器对数目

来源第 125 场双周赛 Q3难度中等分数1908

题目描述

给你一棵无根带权树，树中总共有 n 个节点，分别表示 n 个服务器，服务器从 0 到 n - 1 编号。同时给你一个数组 edges ，其中 edges[i] = [a_i, b_i, weight_i] 表示节点 a_i 和 b_i 之间有一条双向边，边的权值为 weight_i 。再给你一个整数 signalSpeed 。

如果两台服务器 a 和 b 是通过服务器 c 可连接的，则：

a < b ，a != c 且 b != c 。
从 c 到 a 的距离是可以被 signalSpeed 整除的。
从 c 到 b 的距离是可以被 signalSpeed 整除的。
从 c 到 b 的路径与从 c 到 a 的路径没有任何公共边。

请你返回一个长度为 n 的整数数组 count ，其中 count[i] 表示通过服务器 i 可连接 的服务器对的数目。

示例 1：

输入：edges = [[0,1,1],[1,2,5],[2,3,13],[3,4,9],[4,5,2]], signalSpeed = 1
输出：[0,4,6,6,4,0]
解释：由于 signalSpeed 等于 1 ，count[c] 等于所有从 c 开始且没有公共边的路径对数目。
在输入图中，count[c] 等于服务器 c 左边服务器数目乘以右边服务器数目。

示例 2：

输入：edges = [[0,6,3],[6,5,3],[0,3,1],[3,2,7],[3,1,6],[3,4,2]], signalSpeed = 3
输出：[2,0,0,0,0,0,2]
解释：通过服务器 0 ，有 2 个可连接服务器对(4, 5) 和 (4, 6) 。
通过服务器 6 ，有 2 个可连接服务器对 (4, 5) 和 (0, 5) 。
所有服务器对都必须通过服务器 0 或 6 才可连接，所以其他服务器对应的可连接服务器对数目都为 0 。

提示：

2 <= n <= 1000
edges.length == n - 1
edges[i].length == 3
0 <= a_i, b_i < n
edges[i] = [a_i, b_i, weight_i]
1 <= weight_i <= 10⁶
1 <= signalSpeed <= 10⁶
输入保证 edges 构成一棵合法的树。

解法

方法一：枚举 + DFS

思考

以某结点为中转、两条边不相交的路径，两端到中转的距离都整除 \(\textit{signalSpeed}\)。 \(n \le 1000\)。

固定中转 \(a\) 后，不同子树中的合法叶子两两配对。子树内可用 DFS 统计距离整除的结点数。

对每个 \(a\) 枚举儿子，DFS 得到 \(t\)，与此前子树累计 \(s\) 相乘加入答案，再把 \(t\) 累入 \(s\)。

我们先根据题目给定的边构建出一个邻接表 \(g\)，其中 \(g[a]\) 表示节点 \(a\) 的所有邻居节点以及对应的边权。

然后，我们可以枚举每一个节点 \(a\) 作为连接的中间节点，通过深度优先搜索计算出从 \(a\) 的邻居节点 \(b\) 出发的，且到节点 \(a\) 的距离可以被 \(signalSpeed\) 整除的节点数 \(t\)。那么，节点 \(a\) 的可连接节点对数目增加了 \(s \times t\)，其中 \(s\) 表示节点 \(a\) 的邻居节点 \(b\) 出发的，且到节点 \(a\) 的距离不可以被 \(signalSpeed\) 整除的累计节点数。然后我们更新 \(s\) 为 \(s + t\)。

枚举完所有节点 \(a\) 之后，我们就可以得到所有节点的可连接节点对数目。

时间复杂度 \(O(n^2)\)，空间复杂度 \(O(n)\)。其中 \(n\) 表示节点数。

Python3JavaC++GoTypeScript

class Solution:
    def countPairsOfConnectableServers(
        self, edges: List[List[int]], signalSpeed: int
    ) -> List[int]:
        def dfs(a: int, fa: int, ws: int) -> int:
            cnt = 0 if ws % signalSpeed else 1
            for b, w in g[a]:
                if b != fa:
                    cnt += dfs(b, a, ws + w)
            return cnt

        n = len(edges) + 1
        g = [[] for _ in range(n)]
        for a, b, w in edges:
            g[a].append((b, w))
            g[b].append((a, w))
        ans = [0] * n
        for a in range(n):
            s = 0
            for b, w in g[a]:
                t = dfs(b, a, w)
                ans[a] += s * t
                s += t
        return ans

class Solution {
    private int signalSpeed;
    private List<int[]>[] g;

    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
        int n = edges.length + 1;
        g = new List[n];
        this.signalSpeed = signalSpeed;
        Arrays.setAll(g, k -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1], w = e[2];
            g[a].add(new int[] {b, w});
            g[b].add(new int[] {a, w});
        }
        int[] ans = new int[n];
        for (int a = 0; a < n; ++a) {
            int s = 0;
            for (var e : g[a]) {
                int b = e[0], w = e[1];
                int t = dfs(b, a, w);
                ans[a] += s * t;
                s += t;
            }
        }
        return ans;
    }

    private int dfs(int a, int fa, int ws) {
        int cnt = ws % signalSpeed == 0 ? 1 : 0;
        for (var e : g[a]) {
            int b = e[0], w = e[1];
            if (b != fa) {
                cnt += dfs(b, a, ws + w);
            }
        }
        return cnt;
    }
}

class Solution {
public:
    vector<int> countPairsOfConnectableServers(vector<vector<int>>& edges, int signalSpeed) {
        int n = edges.size() + 1;
        vector<pair<int, int>> g[n];
        for (auto& e : edges) {
            int a = e[0], b = e[1], w = e[2];
            g[a].emplace_back(b, w);
            g[b].emplace_back(a, w);
        }
        function<int(int, int, int)> dfs = [&](int a, int fa, int ws) {
            int cnt = ws % signalSpeed == 0;
            for (auto& [b, w] : g[a]) {
                if (b != fa) {
                    cnt += dfs(b, a, ws + w);
                }
            }
            return cnt;
        };
        vector<int> ans(n);
        for (int a = 0; a < n; ++a) {
            int s = 0;
            for (auto& [b, w] : g[a]) {
                int t = dfs(b, a, w);
                ans[a] += s * t;
                s += t;
            }
        }
        return ans;
    }
};

func countPairsOfConnectableServers(edges [][]int, signalSpeed int) []int {
    n := len(edges) + 1
    type pair struct{ x, w int }
    g := make([][]pair, n)
    for _, e := range edges {
        a, b, w := e[0], e[1], e[2]
        g[a] = append(g[a], pair{b, w})
        g[b] = append(g[b], pair{a, w})
    }
    var dfs func(a, fa, ws int) int
    dfs = func(a, fa, ws int) int {
        cnt := 0
        if ws%signalSpeed == 0 {
            cnt++
        }
        for _, e := range g[a] {
            b, w := e.x, e.w
            if b != fa {
                cnt += dfs(b, a, ws+w)
            }
        }
        return cnt
    }
    ans := make([]int, n)
    for a := 0; a < n; a++ {
        s := 0
        for _, e := range g[a] {
            b, w := e.x, e.w
            t := dfs(b, a, w)
            ans[a] += s * t
            s += t
        }
    }
    return ans
}

function countPairsOfConnectableServers(edges: number[][], signalSpeed: number): number[] {
    const n = edges.length + 1;
    const g: [number, number][][] = Array.from({ length: n }, () => []);
    for (const [a, b, w] of edges) {
        g[a].push([b, w]);
        g[b].push([a, w]);
    }
    const dfs = (a: number, fa: number, ws: number): number => {
        let cnt = ws % signalSpeed === 0 ? 1 : 0;
        for (const [b, w] of g[a]) {
            if (b != fa) {
                cnt += dfs(b, a, ws + w);
            }
        }
        return cnt;
    };
    const ans: number[] = Array(n).fill(0);
    for (let a = 0; a < n; ++a) {
        let s = 0;
        for (const [b, w] of g[a]) {
            const t = dfs(b, a, w);
            ans[a] += s * t;
            s += t;
        }
    }
    return ans;
}

方法二：枚举 + 显式栈

思考

一对服务器能通过中转点 \(a\) 相连，当且仅当两条路径没有公共边，并且两端到 \(a\) 的距离都能被 \(\textit{signalSpeed}\) 整除。\(n \le 1000\)，固定中转点之后，不同分支里的合法节点可以两两配对。每个分支里，距离能被整除的节点个数与访问顺序无关，用一次遍历就能数出来。沿一条链递归时调用深度等于节点数，在 \(n = 1000\) 时会超出 Python 的递归上限。因此栈中保存 \((节点, 父节点, 已走距离)\)：弹出时若距离能被整除就计入，再把其余邻接点以距离加上边权压入。对中转点 \(a\) 的每个邻居得到个数 \(t\)，与此前分支的累计 \(s\) 相乘加入答案，再把 \(t\) 累入 \(s\)。

先根据给定的边建立邻接表 \(g\)，\(g[a]\) 保存节点 \(a\) 的邻居以及对应边权。

枚举每个节点 \(a\) 作为中转点。对 \(a\) 的每个邻居 \(b\)，用显式栈统计 \(b\) 所在分支中、到 \(a\) 的距离能被 \(\textit{signalSpeed}\) 整除的节点数 \(t\)。栈里每个元素是 \((节点, 父节点, 已走距离)\)。从 \(b\) 出发时距离就是边权，弹出一个节点后，若当前距离能被整除就把计数加一，再把除父节点以外的邻接点压入，距离加上对应边权。记此前各分支已经统计到的合法节点数为 \(s\)，则这些节点都能与当前分支的 \(t\) 个节点配对，答案增加 \(s \times t\)，然后把 \(t\) 累加进 \(s\)。分支的处理顺序不影响配对数。

枚举完所有中转点后，就得到每个节点作为中转时的可连接对数。