Find number of coins to place in tree nodes

Dynamic Programming Tree Depth-First Search Sorting Heap (Priority Queue)

37.3% acceptance

Feb 25, 2026

201

You are given an undirected tree rooted at node 0. You are given edges and a cost array. Place coins at each node: if subtree size < 3, place 1 coin; else place max(0, max product of 3 costs in subtree). Return array coin[i] = coins at node i.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn placed_coins(edges: Vec<Vec<i32>>, cost: Vec<i32>) -> Vec<i64> {
    let n = cost.len();
    let mut adj: Vec<Vec<usize>> = vec![vec![]; n];
    for e in &edges {
      let u = e[0] as usize;
      let v = e[1] as usize;
      adj[u].push(v);
      adj[v].push(u);
    }

    let mut coins = vec![0i64; n];

    // DFS: return top-5 values (sorted desc) from subtree
    // BFS/iterative DFS to avoid stack overflow
    fn dfs(node: usize, parent: usize, adj: &Vec<Vec<usize>>, cost: &Vec<i32>, coins: &mut Vec<i64>) -> Vec<i32> {
      let mut vals = vec![cost[node]];
      for &child in &adj[node] {
        if child == parent { continue; }
        let child_vals = dfs(child, node, adj, cost, coins);
        vals.extend(child_vals);
      }
      // Compute coins for this node
      if vals.len() < 3 {
        coins[node] = 1;
      } else {
        vals.sort_unstable_by(|a, b| b.cmp(a)); // descending
        // Option 1: top 3
        let p1 = vals[0] as i64 * vals[1] as i64 * vals[2] as i64;
        // Option 2: bottom 2 (most negative) × top 1
        let m = vals.len();
        let p2 = vals[0] as i64 * vals[m - 1] as i64 * vals[m - 2] as i64;
        coins[node] = 0i64.max(p1.max(p2));
      }
      // Keep only top 3 and bottom 2 for parent
      let m = vals.len();
      if m <= 5 {
        vals
      } else {
        vec![vals[0], vals[1], vals[2], vals[m - 2], vals[m - 1]]
      }
    }

    dfs(0, n, &adj, &cost, &mut coins);
    coins
  }
}