LC每日一题｜20240604 - 3067. 在带权树网络中统计可连接服务器对数目

给你一棵无根带权树，树中总共有 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 可连接 的服务器对的 数目 。

提示：

• 2 <= n <= 1000
• edges.length == n - 1
• edges[i].length == 3
• 0 <= ai, bi < n
• edges[i] = [ai, bi, weighti]
• 1 <= weighti <= 10^6
• 1 <= signalSpeed <= 10^6
• 输入保证 edges 构成一棵合法的树。

题目等级：Medium

解题思路

一道不是很难的dfs题～

我们可以先建树，然后枚举所有的根，再dfs统计这个根的每个树枝上满足条件的节点的个数，然后累加各枝到该节点下其他枝的节点数量的乘积即可。

AC代码

class Solution {
    fun countPairsOfConnectableServers(edges: Array<IntArray>, k: Int): IntArray {
        val res = IntArray(edges.size + 1)
        val map = Array<ArrayList<IntArray>>(edges.size + 1) { arrayListOf() }
        edges.forEach {
            map[it[0]].add(intArrayOf(it[1], it[2]))
            map[it[1]].add(intArrayOf(it[0], it[2]))
        }

        fun dfs(cur: Int, dis: Int, parent: Int): Int {
            var res = 0
            if (dis % k == 0) res += 1
            map[cur].forEach { 
                if (parent != it[0]) {
                    res += dfs(it[0], dis + it[1], cur)
                }
            }
            return res
        }

        repeat(edges.size + 1) { node ->
            val list = ArrayList<Int>()
            var total = 0
            map[node].forEach {
                val x = dfs(it[0], it[1], node)
                list.add(x)
                total += x
            }
            for (n in list) res[node] += n * (total - n)
            res[node] /= 2
        }
        return res
    }
}