Total sum of interaction cost in tree groups

Array Tree Depth-First Search

53.6% acceptance

Feb 25, 2026

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] = [ui, vi] indicates an undirected edge between nodes ui and vi. 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].

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn interaction_costs(n: i32, edges: Vec<Vec<i32>>, group: Vec<i32>) -> i64 {
    let n = n as usize;
    let mut adj = vec![vec![]; n];
    for e in &edges {
      let (u, v) = (e[0] as usize, e[1] as usize);
      adj[u].push(v);
      adj[v].push(u);
    }
    // Count total per group
    let max_g = *group.iter().max().unwrap_or(&0) as usize;
    let mut total_count = vec![0i64; max_g + 1];
    for &g in &group { total_count[g as usize] += 1; }

    // DFS to compute subtree group counts
    let mut subtree = vec![vec![0i64; max_g + 1]; n];
    for i in 0..n { subtree[i][group[i] as usize] = 1; }

    // BFS order with parent tracking
    let mut parent = vec![n; n];
    let mut order = Vec::with_capacity(n);
    let mut visited = vec![false; n];
    let mut queue = std::collections::VecDeque::new();
    queue.push_back(0usize);
    visited[0] = true;
    while let Some(u) = queue.pop_front() {
      order.push(u);
      for &v in &adj[u] {
        if !visited[v] {
          visited[v] = true;
          parent[v] = u;
          queue.push_back(v);
        }
      }
    }
    // Bottom-up: accumulate subtree counts
    for &u in order.iter().rev() {
      let p = parent[u];
      if p < n {
        for g in 0..=max_g {
          let c = subtree[u][g];
          subtree[p][g] += c;
        }
      }
    }
    // For each edge (parent[v], v): contribution = sum_g sub[v][g] * (total[g] - sub[v][g])
    let mut ans = 0i64;
    for &v in &order {
      if v == 0 { continue; }
      for g in 0..=max_g {
        let left = subtree[v][g];
        let right = total_count[g] - left;
        ans += left * right;
      }
    }
    ans
  }
}