Learn AI Create the Future — with Pranay
Main site

LeetCode solutions

4072. Total Sum of Interaction Cost in Tree Groups

My accepted Python solution to LeetCode problem 4072, Total Sum of Interaction Cost in Tree Groups, running in 795ms.

  • Difficulty: Hard
  • Python
  • Runtime 795ms
  • Memory 113.2MB

Read the problem on LeetCode View on GitHub

The problem statement is LeetCode’s and stays on their site. What follows is my accepted solution.

Python

Accepted on LeetCode — runtime 795ms, memory 113.2MB, accepted 2025-12-29.

python
class Solution:
    def interactionCosts(self, n: int, edges: List[List[int]], group: List[int]) -> int:
        from collections import defaultdict
        adj = defaultdict(list)
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)
        group_count = defaultdict(int)
        for g in group:
            group_count[g] += 1
        self.result = 0
        def dfs(node, parent):
            subtree_count = defaultdict(int)
            subtree_count[group[node]] = 1
            for neighbor in adj[node]:
                if neighbor != parent:
                    child_count = dfs(neighbor, node)
                    for g in child_count:
                        if group_count[g] > 0:
                            self.result += child_count[g] * (group_count[g] - child_count[g])
                    for g, cnt in child_count.items():
                        subtree_count[g] += cnt
            return subtree_count
        dfs(0, -1)
        return self.result

Source