Collect Coins in a Tree

hard tree topological sort bfs

Problem

An undirected tree has n nodes given by an edge list, and coins[i] is 1 if node i has a coin. Starting from any node you may move along edges; from a node you can collect every coin within distance 2, and you must return to your start. Return the minimum number of edges you need to traverse (each traversed edge is counted twice for the round trip).

Inputcoins = [1,0,0,0,0,1], edges = [[0,1],[1,2],[2,3],[3,4],[4,5]]

Output2

After trimming, two nodes remain, so you walk a single edge there and back: 1 edge × 2 = 2.

coins array

edges as [u,v] list

def collect_the_coins(coins, edges):
    n = len(coins)
    g = [set() for _ in range(n)]
    for a, b in edges:
        g[a].add(b); g[b].add(a)
    deg = [len(g[i]) for i in range(n)]
    removed = [False] * n
    remaining = n
    leaves = [i for i in range(n) if deg[i] <= 1 and coins[i] == 0]
    while leaves:
        nxt = []
        for u in leaves:
            removed[u] = True; remaining -= 1
            for v in g[u]:
                g[v].discard(u); deg[v] -= 1
                if deg[v] == 1 and coins[v] == 0:
                    nxt.append(v)
            g[u].clear()
        leaves = nxt
    for _ in range(2):
        leaves = [i for i in range(n) if not removed[i] and deg[i] == 1]
        for u in leaves:
            removed[u] = True; remaining -= 1
            for v in g[u]:
                deg[v] -= 1
            g[u].clear()
    return max(remaining - 1, 0) * 2

function collectTheCoins(coins, edges) {
  const n = coins.length;
  const g = Array.from({length: n}, () => new Set());
  const deg = new Array(n).fill(0);
  for (const [a, b] of edges) { g[a].add(b); g[b].add(a); deg[a]++; deg[b]++; }
  const removed = new Array(n).fill(false);
  let remaining = n;
  let leaves = [];
  for (let i = 0; i < n; i++) if (deg[i] <= 1 && coins[i] === 0) leaves.push(i);
  while (leaves.length) {
    const nxt = [];
    for (const u of leaves) {
      removed[u] = true; remaining--;
      for (const v of g[u]) {
        g[v].delete(u); deg[v]--;
        if (deg[v] === 1 && coins[v] === 0) nxt.push(v);
      }
      g[u].clear();
    }
    leaves = nxt;
  }
  for (let r = 0; r < 2; r++) {
    const cur = [];
    for (let i = 0; i < n; i++) if (!removed[i] && deg[i] === 1) cur.push(i);
    for (const u of cur) {
      removed[u] = true; remaining--;
      for (const v of g[u]) deg[v]--;
      g[u].clear();
    }
  }
  return Math.max(remaining - 1, 0) * 2;
}

import java.util.*;
class Solution {
    public int collectTheCoins(int[] coins, int[][] edges) {
        int n = coins.length;
        Set<Integer>[] g = new HashSet[n];
        for (int i = 0; i < n; i++) g[i] = new HashSet<>();
        int[] deg = new int[n];
        for (int[] e : edges) { g[e[0]].add(e[1]); g[e[1]].add(e[0]); deg[e[0]]++; deg[e[1]]++; }
        boolean[] removed = new boolean[n];
        int remaining = n;
        Queue<Integer> q = new LinkedList<>();
        for (int i = 0; i < n; i++) if (deg[i] <= 1 && coins[i] == 0) q.add(i);
        while (!q.isEmpty()) {
            int u = q.poll(); removed[u] = true; remaining--;
            for (int v : g[u]) { g[v].remove(u); if (--deg[v] == 1 && coins[v] == 0) q.add(v); }
            g[u].clear();
        }
        for (int r = 0; r < 2; r++) {
            List<Integer> cur = new ArrayList<>();
            for (int i = 0; i < n; i++) if (!removed[i] && deg[i] == 1) cur.add(i);
            for (int u : cur) { removed[u] = true; remaining--; for (int v : g[u]) deg[v]--; g[u].clear(); }
        }
        return Math.max(remaining - 1, 0) * 2;
    }
}

class Solution {
public:
    int collectTheCoins(vector<int>& coins, vector<vector<int>>& edges) {
        int n = coins.size();
        vector<set<int>> g(n);
        vector<int> deg(n, 0);
        for (auto& e : edges) { g[e[0]].insert(e[1]); g[e[1]].insert(e[0]); deg[e[0]]++; deg[e[1]]++; }
        vector<bool> removed(n, false);
        int remaining = n;
        queue<int> q;
        for (int i = 0; i < n; i++) if (deg[i] <= 1 && coins[i] == 0) q.push(i);
        while (!q.empty()) {
            int u = q.front(); q.pop(); removed[u] = true; remaining--;
            for (int v : g[u]) { g[v].erase(u); if (--deg[v] == 1 && coins[v] == 0) q.push(v); }
            g[u].clear();
        }
        for (int r = 0; r < 2; r++) {
            vector<int> cur;
            for (int i = 0; i < n; i++) if (!removed[i] && deg[i] == 1) cur.push_back(i);
            for (int u : cur) { removed[u] = true; remaining--; for (int v : g[u]) deg[v]--; g[u].clear(); }
        }
        return max(remaining - 1, 0) * 2;
    }
};

Explanation

Because you collect coins within distance 2 and must return to start, the answer only depends on which nodes you must actually walk through. Three rounds of leaf trimming reduce the tree to exactly that core.

First, repeatedly remove every coinless leaf. A coinless leaf is a dead end with nothing to collect, and removing it can expose new coinless leaves, so we keep peeling (a topological / multi-source BFS). After this, every remaining leaf carries a coin.

Second, trim the leaves exactly twice. Standing two steps before a coin-leaf already lets you collect it, so you never need to walk all the way out to a coin or to its neighbor. Two passes shave off those last two layers.

Whatever nodes survive form the set you genuinely must traverse. If remaining nodes are left, the minimal walk that visits them all and returns uses (remaining - 1) edges, each crossed twice, for an answer of 2 · (remaining - 1). If nothing is left (no coins or a tiny tree), the answer is 0.

Example on the path 0–1–2–3–4–5 with coins at 0 and 5: there are no coinless leaves to peel, then two leaf-trims remove nodes 0,5 and then 1,4, leaving {2,3}. That is (2 - 1) · 2 = 2.

Time: O(n) Space: O(n)