Leetcode 3067. Count Pairs of Connectable Servers in a Weighted Tree Network

• Leetcode 3067. Count Pairs of Connectable Servers in a Weighted Tree Network
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：3067. Count Pairs of Connectable Servers in a Weighted Tree Network

1. 解题思路

这一题没想到什么好的方法，走的是暴力求解的路子。

整体的思路上其实还是比较直接的，就是考察所有的节点作为根节点时的单项路径，然后筛选出其中所有的和可以被signalSpeed整除的路径数目，最后求一下两两组合的对数即可。

其中，前者可以通过一个深度优先遍历进行实现；而后者事实上就是如下公式：

$$\begin{array}{rl}s& =\sum _{i<j}{x}_i\cdot x_j\\ & =\frac{1}{2}\cdot \sum _i{x}_i\cdot (\sum _{j\ne i}{x}_j)\\ & =\frac{1}{2}\cdot \sum _i{x}_i\cdot (\sum _j{x}_j-x_i)\end{array}$$

因此，我们通过一重循环就能够快速地得到对应的答案。

2. 代码实现

给出python代码实现如下：

class Solution:
    def countPairsOfConnectableServers(self, edges: List[List[int]], signalSpeed: int) -> List[int]:
        n = len(edges)+1
        graph = defaultdict(list)
        for u, v, w in edges:
            graph[u].append((v, w))
            graph[v].append((u, w))

        @lru_cache(None)
        def dfs(u, pre, distance):
            ans =  1 if distance == 0 else 0
            for v, w in graph[u]:
                if v == pre:
                    continue
                ans += dfs(v, u, (distance + w) % signalSpeed)
            return ans
        
        def count_connectable(u):
            cnt = [dfs(v, u, w % signalSpeed) for v, w in graph[u]]
            s = sum(cnt)
            ans = sum([x * (s-x) for x in cnt]) // 2
            return ans
        
        return [count_connectable(u) for u in range(n)] 

提交代码评测得到：耗时292ms，占用内存38.7MB。