【Leetcode】3067-1909

前言

去年做过的题目，今天竟然依旧做了好久。

---

思路

因为题目需求返回的是每个节点作为中间服务器的，因此需要每个节点分别的进行计算。

考虑其中的一个中间节点，遍历其中的每个分支可能的连接，再乘法累加原理就得到了每个节点可能的服务器对数了。

什么是乘法累加原理呢？考虑a,b,c,d等若干个集合中分别选择两个集合中的某一元素进行组合，计算有多少个组合。

按照最简单的做法就是枚举每两个集合，分别计算每个集合对可能的组合，算法的时间复杂度为\(O(n^2)\)，其中\(n\)为集合的数目

而针对于乘法累加原理而言，我们可以考虑一次遍历，同时记录当前已经遍历过的所有集合中元素的总数。每一次遍历中，将当前集合的元素个数乘以记录的所有种元素的总数即可。该算法的时间复杂度为\(O(n)\).

class Solution:
    def countPairsOfConnectableServers(self, edges: List[List[int]], s: int) -> List[int]:
        n = len(edges)+1
        g = [[] for _ in range(n)]
        for u, v, w in edges:
            g[u].append((v, w))
            g[v].append((u, w))
        def dfs(u, fa, cur):
            res =int( cur==0)
            for v, w in g[u]:
                if v == fa: continue
                res += dfs(v, u, (cur + s - w) % s)
            return res
        ans = [0] * n
        for u in range(n ):
            cur = 0
            if len(g[u])==1:continue
            for v, w in g[u]:
                temp = dfs(v, u, (s - w % s)%s)
                ans[u] += temp * cur
                cur += temp
        return ans

上述代码中，dfs(u,fa,cur)表示考虑以fa为父节点，当前节点为u，查找到距离与s的余数为cur的路线的个数

当然，考虑到实际上有很多重复的计算，也是可以直接将dfs(u,fa,cur)的结果进行存储。不过与此同时也换来了更多的内存空间使用。

@cache
def dfs(u, fa, cur):
    res =int( cur==0)
    for v, w in g[u]:
        if v == fa: continue
        res += dfs(v, u, (cur + s - w) % s)
    return res