3786. Total Sum of Interaction Cost in Tree Groups

You are given an integer n and an undirected tree with n nodes numbered from 0 to n - 1. This is represented by a 2D array edges of length n - 1, where edges[i] = [u_i, v_i] indicates an undirected edge between nodes u_i and v_i.

You are also given an integer array group of length n, where group[i] denotes the group label assigned to node i.

• Two nodes u and v are considered part of the same group if group[u] == group[v].
• The interaction cost between u and v is defined as the number of edges on the unique path connecting them in the tree.

Return an integer denoting the sum of interaction costs over all unordered pairs (u, v) with u != v such that group[u] == group[v].

class Solution {
    unordered_map<int,pair<long long,long long>> dfs(int u, int par, vector<vector<int>>& adj, vector<int>& g, long long& res) {
        unordered_map<int,pair<long long,long long>> acc;
        for(auto& v : adj[u]) {
            if(v == par) continue;
            auto sub = dfs(v,u,adj,g,res);
            for(auto& [k,p] : sub) {
                if(k == g[u]) res += p.first;
                if(acc.contains(k)) {
                    res += acc[k].second * p.first;
                    res += p.second * acc[k].first;
                }
                acc[k].first += p.first;
                acc[k].second += p.second;
            }
        }
        acc[g[u]].second += 1;
        for(auto& [_,p] : acc) {
            p.first += p.second;
        }
        return acc;
    }
public:
    long long interactionCosts(int n, vector<vector<int>>& edges, vector<int>& group) {
        vector<vector<int>> adj(n);
        for(auto& e : edges) {
            int u = e[0], v = e[1];
            adj[u].push_back(v);
            adj[v].push_back(u);
        }
        long long res = 0;
        dfs(0,-1,adj,group,res);
        return res;
    }
};