Leetcode 3786. Total Sum of Interaction Cost in Tree Groups

• Leetcode 3786. Total Sum of Interaction Cost in Tree Groups
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：3786. Total Sum of Interaction Cost in Tree Groups

1. 解题思路

这一题是Leetcode周赛481的第四题，是一道hard难度的题目。

这一题我没有自己搞定，是看了一下大佬的解题方法搞定的。

这一题的核心是需要转换一下思路，不是两两去考虑任意两个点的距离，而是去考虑任意一条边会被多少条满足条件的路径使用到。

此时，我们就可以通过图遍历的方式考察任意一条边下游的子树当中每一个group各自有多少个节点，此时，剩余所有的属于该group的节点要与之相连的话都需要经过这条边。

由此，我们就可以给出python代码实现了。

2. 代码实现

给出python代码实现如下：

class Solution:
    def interactionCosts(self, n: int, edges: List[List[int]], group: List[int]) -> int:
        graph = defaultdict(list)
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        total = Counter(group)
        ans = 0
        
        def dfs(u, fa):
            nonlocal ans
            cnt = defaultdict(int)
            cnt[group[u]] = 1
            for v in graph[u]:
                if v == fa:
                    continue
                subtree = dfs(v, u)
                for grp, _cnt in subtree.items():
                    ans += _cnt * (total[grp] - _cnt)
                    cnt[grp] += _cnt
            return cnt

        dfs(0, -1)
        return ans

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