Day 12：3067. 在带权树网络中统计可连接服务器对数目-CSDN博客

本文链接：https://blog.csdn.net/weixin_46091073/article/details/139435545

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

如果两个服务器 a ，b 和 c 满足以下条件，那么我们称服务器 a 和 b 是通过服务器 c 可连接的 ：

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

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

选择 c 作为根节点，不同子树中进行深度搜索，在不同子树的节点，就不会有公共边，不同子树中满足条件的节点数相乘就是所得结果。

如何存储其连接和权值，可以使用List<int[]>[]进行存储，也可以使用二维数组进行存储。

不需要考虑 a 和 b 的大小关系，我们选择一对，总有一个大的值。

还有就是条件并没有说明，它是一棵无根二叉树，它可能是多叉树。因此要计算某一棵子树与其他剩余子树满足条件的节点。

完整代码

class Solution {
    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
        int n = edges.length + 1;

        List<int[]>[] graph = new ArrayList[n + 1];
        for (int i = 0; i < n + 1; i++) {
            graph[i] = new ArrayList<>();
        }
        for (int[] edge : edges) {
            int u = edge[0];
            int v = edge[1];
            int w = edge[2];
            graph[u].add(new int[]{v, w});
            graph[v].add(new int[]{u, w});
        }

        int[] res = new int[n];
        for (int i = 0; i < n; i++) {
            int p = 0;
            for (int[] num : graph[i]) {
                int cnt = dfs(i, num[0], num[1], signalSpeed, graph);
                res[i] += p * cnt;
                p += cnt;
            }
        }
        return res;
    }

    public int dfs(int pre, int cur, int cost, int signalSpeed, List<int[]>[] graph) {
        int res = 0;
        if (cost % signalSpeed == 0) {
            res++;
        }
        for (int[] num : graph[cur]) {
            if (num[0] == pre) continue;
            res += dfs(cur, num[0], cost + num[1], signalSpeed, graph);
        }
        return res;
    }
}

用二维数组存储的完整代码

class Solution {
    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
        int n = edges.length + 1;

        int[][] weight = new int[n][n];
        for (int[] edge : edges) {
            int u = edge[0];
            int v = edge[1];
            int w = edge[2];
            weight[u][v] = w;
            weight[v][u] = w;
        }

        int[] res = new int[n];
        for (int i = 0; i < n; i++) {
            int p = 0;
            for (int j = 0; j < n; j++) {
                if (weight[i][j] == 0 || j == i) continue;
                int cnt = dfs(n, i, j, weight[i][j], signalSpeed, weight);
                res[i] += p * cnt;
                p += cnt;
            }
        }
        return res;
    }

    public int dfs(int n, int pre, int cur, int cost, int signalSpeed, int[][] weight) {
        int res = 0;
        if (cost % signalSpeed == 0) {
            res++;
        }
        for (int j = 0; j < n; j++) {
            if (weight[cur][j] == 0 || j == cur || j == pre) continue;
            res += dfs(n, cur, j, cost + weight[cur][j], signalSpeed, weight);
        }
        return res;
    }
}