Collect coins in a tree

Array Tree Graph Theory Topological Sort

39.7% acceptance

Feb 25, 2026

558

There exists an undirected and unrooted tree with n nodes indexed from 0 to n - 1. You are given an integer n and a 2D integer array edges of length n - 1, where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree. You are also given an array coins of size n where coins[i] can be either 0 or 1, where 1 indicates the presence of a coin in the vertex i. Initially, you choose to start at any vertex in the tree. Then, you can perform the following operations any number of times: Collect all the coins that are at a distance of at most 2 from the current vertex, or Move to any adjacent vertex in the tree. Find the minimum number of edges you need to go through to collect all the coins and go back to the initial vertex. Note that if you pass an edge several times, you need to count it into the answer several times.

Solution

Rust

Time O(n³)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn collect_the_coins(coins: Vec<i32>, edges: Vec<Vec<i32>>) -> i32 {
    let n = coins.len();
    if n <= 2 {
      return 0;
    }
    let mut adj: Vec<std::collections::HashSet<usize>> = vec![std::collections::HashSet::new(); n];
    for e in &edges {
      let (a, b) = (e[0] as usize, e[1] as usize);
      adj[a].insert(b);
      adj[b].insert(a);
    }
    let mut degree: Vec<usize> = (0..n).map(|i| adj[i].len()).collect();
    let mut removed = vec![false; n];

    // Step 1: remove leaves with no coins repeatedly
    let mut queue: std::collections::VecDeque<usize> = std::collections::VecDeque::new();
    for i in 0..n {
      if degree[i] == 1 && coins[i] == 0 {
        queue.push_back(i);
      }
    }
    while let Some(u) = queue.pop_front() {
      removed[u] = true;
      for &v in &adj[u].clone() {
        adj[v].remove(&u);
        degree[v] -= 1;
        if degree[v] == 1 && coins[v] == 0 && !removed[v] {
          queue.push_back(v);
        }
      }
    }

    // Step 2: remove leaves twice more (distance-2 collection radius)
    for _ in 0..2 {
      let leaves: Vec<usize> = (0..n).filter(|&i| !removed[i] && adj[i].len() <= 1).collect();
      for u in leaves {
        removed[u] = true;
        for &v in &adj[u].clone() {
          adj[v].remove(&u);
        }
        adj[u].clear();
      }
    }

    // Count remaining edges
    let remaining_nodes: usize = (0..n).filter(|&i| !removed[i]).count();
    if remaining_nodes <= 1 {
      return 0;
    }
    // remaining edges = remaining_nodes - 1 (tree), answer = 2 * edges
    2 * (remaining_nodes as i32 - 1)
  }
}