# Collect Coins in a Tree
**Difficulty:** HARD
[External](https://leetcode.com/problems/collect-coins-in-a-tree)
Canonical: https://scaleengineer.com/dsa/problems/collect-coins-in-a-tree
**Algorithms:** [Topological Sort](https://scaleengineer.com/algorithms/topological-sort)
**Data structures:** Array, Tree, Graph
**Companies:** [Lucid Motors](https://scaleengineer.com/companies/lucid-motors), [Graviton](https://scaleengineer.com/companies/graviton)
---
## Problem
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.

**Example 1:**

![](https://assets.glich.co/dsa/collect-coins-in-a-tree/image0.png) 

**Input:** coins = [1,0,0,0,0,1], edges = [[0,1],[1,2],[2,3],[3,4],[4,5]]
**Output:** 2
**Explanation:** Start at vertex 2, collect the coin at vertex 0, move to vertex 3, collect the coin at vertex 5 then move back to vertex 2.

**Example 2:**

![](https://assets.glich.co/dsa/collect-coins-in-a-tree/image1.png) 

**Input:** coins = [0,0,0,1,1,0,0,1], edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[5,6],[5,7]]
**Output:** 2
**Explanation:** Start at vertex 0, collect the coins at vertices 4 and 3, move to vertex 2,  collect the coin at vertex 7, then move back to vertex 0.

**Constraints:**

* `n == coins.length`
* `1 <= n <= 3 * 104`
* `0 <= coins[i] <= 1`
* `edges.length == n - 1`
* `edges[i].length == 2`
* `0 <= ai, bi < n`
* `ai != bi`
* `edges` represents a valid tree.

# Approaches
## Iterative Tree Pruning
This approach is based on the idea of progressively trimming the tree to find the minimal subtree that must be traversed. The total travel distance is twice the number of edges in this minimal subtree. The pruning happens in two main stages. First, we remove all subtrees that contain no coins. Second, we account for the collection rule (distance at most 2) by trimming two additional layers of nodes from the leaves of the remaining tree. This version implements the pruning using simple loops, which can be less efficient on certain tree structures.
**Time:** O(N^2). In the worst case, such as a path graph, each pass of the `while` loop might only prune one node. Since there can be O(N) nodes to prune in total, the complexity becomes O(N * N). · **Space:** O(N) to store the adjacency list, an array for node degrees, and a boolean array to track pruned nodes.
**Pros:** The logic is broken down into simple, sequential steps, which can be easier to reason about initially.
**Cons:** The nested loop structure for finding and removing leaves one by one can be very slow, leading to a quadratic time complexity in the worst-case scenarios, such as a path-like tree.; It is more complex to implement correctly compared to a queue-based approach, as it requires careful state management (e.g., `isPruned` array) and repeated iterations.
### Explanation
The fundamental insight is that we only need to traverse a minimal subtree that allows us to be within a distance of 2 from every coin. Any parts of the tree that are not required for this can be pruned.

This method implements the pruning iteratively. It repeatedly scans all nodes to find and remove leaves that are not needed. A leaf node `u` with `coins[u] == 0` can be removed because it's not a destination and isn't needed to connect to other coins. After removing all such subtrees, we are left with a 'core tree'. For this core tree, we don't need to visit the leaf nodes themselves or their parents to collect the coins at the leaves, thanks to the distance-2 collection rule. Therefore, we can prune two layers of leaves from the core tree. The iterative implementation of this pruning, while conceptually straightforward, leads to poor performance.

```java
// Illustrative pseudocode for the iterative pruning logic

// Stage 1: Prune useless subtrees
boolean changed = true;
while (changed) {
    changed = false;
    for (int i = 0; i < n; i++) {
        // Check if node i is a leaf and can be pruned
        if (!isPruned[i] && degree[i] == 1 && coins[i] == 0) {
            // Prune node i
            isPruned[i] = true;
            degree[i]--;
            // Find neighbor v and decrement degree[v]
            for (int neighbor : adj.get(i)) {
                if (!isPruned[neighbor]) {
                    degree[neighbor]--;
                }
            }
            changed = true;
        }
    }
}

// Stage 2: Prune two layers for collection distance
// This would involve two more similar while loops, one for each layer.
```
### Algorithm
- Build an adjacency list representation of the tree and calculate the initial degree of each node.
- **Stage 1: Prune subtrees with no coins.**
  - Enter a loop that continues as long as nodes are being pruned.
  - Inside the loop, iterate through all nodes from `0` to `n-1`.
  - If a node `i` is a leaf (degree 1), has not been pruned, and has no coin (`coins[i] == 0`), prune it.
  - Pruning involves marking the node as removed, decrementing its degree, and decrementing the degree of its single neighbor.
  - The loop terminates when a full pass over all nodes prunes nothing.
- **Stage 2: Prune for collection distance.**
  - After Stage 1, the remaining nodes form the smallest subtree connecting all coin-bearing nodes (the "core tree").
  - Apply a similar iterative pruning process twice to remove two layers of leaves from this core tree.
  - **Layer 1:** In a loop, find and prune all leaves of the current tree until no more leaves can be pruned.
  - **Layer 2:** Repeat the process to prune the new layer of leaves that emerged after the first layer's removal.
- **Calculate Result:**
  - Count the number of nodes `R` that were not pruned.
  - The number of edges in the final required traversal path is `R - 1`.
  - The total distance is `2 * (R - 1)`. If `R <= 1`, the cost is 0.

## Optimized Pruning with Queues
This is an optimized version of the pruning approach. Instead of repeatedly scanning all nodes to find leaves to prune, it uses queues to efficiently process nodes in a way similar to topological sorting or Breadth-First Search (BFS). This allows for pruning the tree in linear time. The core logic remains the same: first, trim subtrees without coins, then trim two layers to account for the collection distance. This is the most efficient way to solve the problem.
**Time:** O(N). Building the graph takes O(N). Each pruning stage involves processing each node and edge at most once, similar to a BFS or topological sort. Therefore, the total time complexity is linear. · **Space:** O(N). We need an adjacency structure (list or set) to represent the tree, which takes O(N) space for N-1 edges. The queues used for pruning will also take at most O(N) space.
**Pros:** Highly efficient with a linear time complexity, making it suitable for large inputs.; Provides a direct and elegant solution by reducing the problem to its essential core through pruning.; The use of queues streamlines the pruning process, making the implementation clean and avoiding nested loops.
**Cons:** The logic, especially the two-layer pruning and its interaction with the graph representation, can be subtle to implement correctly without errors.; Requires careful management of the graph state as nodes and edges are effectively removed.
### Explanation
This approach refines the iterative pruning method by using queues for efficient, single-pass removal of nodes. This avoids the O(N^2) complexity and solves the problem in linear time.

**Stage 1:** We perform a topological sort-like removal of all leaf nodes that do not contain coins. A queue is initialized with these leaves. As we process the queue, a node is removed, and its neighbor's degree (or adjacency list size) is updated. If the neighbor becomes a leaf with no coin, it's added to the queue. This efficiently strips away all irrelevant subtrees.

**Stage 2:** After the first stage, we are left with the core subtree. We then need to trim two layers of leaves from this subtree. This is done by first identifying all current leaves, removing them, and then identifying the next set of leaves and removing them. This process efficiently finds the final minimal set of nodes that must be part of the traversal path.

The final answer is twice the number of edges in the remaining subtree.

```java
import java.util.*;

class Solution {
    public int collectCoins(int[] coins, int[][] edges) {
        int n = coins.length;
        if (n <= 1) {
            return 0;
        }

        Set<Integer>[] adj = new HashSet[n];
        for (int i = 0; i < n; i++) {
            adj[i] = new HashSet<>();
        }
        for (int[] edge : edges) {
            adj[edge[0]].add(edge[1]);
            adj[edge[1]].add(edge[0]);
        }

        // Stage 1: Prune leaves with no coins
        Queue<Integer> q = new LinkedList<>();
        for (int i = 0; i < n; i++) {
            if (adj[i].size() == 1 && coins[i] == 0) {
                q.add(i);
            }
        }

        int remainingNodes = n;
        while (!q.isEmpty()) {
            int u = q.poll();
            remainingNodes--;
            if (adj[u].isEmpty()) continue;
            int v = adj[u].iterator().next();
            adj[v].remove(u);
            if (adj[v].size() == 1 && coins[v] == 0) {
                q.add(v);
            }
        }

        // Stage 2: Prune two layers of all leaves
        for (int i = 0; i < 2; i++) {
            List<Integer> leaves = new ArrayList<>();
            for (int j = 0; j < n; j++) {
                if (adj[j].size() == 1) {
                    leaves.add(j);
                }
            }
            if (leaves.isEmpty()) break;
            for (int leaf : leaves) {
                remainingNodes--;
                if (adj[leaf].isEmpty()) continue;
                int neighbor = adj[leaf].iterator().next();
                adj[neighbor].remove(leaf);
                adj[leaf].clear(); // Mark as removed
            }
        }

        return Math.max(0, (remainingNodes - 1) * 2);
    }
}
```
### Algorithm
- **Graph Representation:** Build an adjacency list (or set) for the tree. Using a `Set` for neighbors makes edge removal O(1) on average.
- **Stage 1: Prune subtrees with no coins.**
  - Create a queue and initialize it with all leaf nodes `u` (nodes with one neighbor) where `coins[u] == 0`.
  - Use a variable `remainingNodes` initialized to `n`.
  - While the queue is not empty, dequeue a node `u`. Decrement `remainingNodes`.
  - Remove the edge connecting `u` to its neighbor `v`. 
  - If `v` becomes a new leaf (`adj[v].size() == 1`) and `coins[v] == 0`, add `v` to the queue.
- **Stage 2: Prune two layers for collection distance.**
  - This stage is performed on the remaining graph.
  - Run a loop twice (for two layers).
  - In each iteration, find all current leaves (nodes `j` with `adj[j].size() == 1`) and store them in a temporary list.
  - If there are no leaves, the graph is empty or a single node, so break.
  - Iterate through the list of leaves. For each leaf, decrement `remainingNodes` and remove it and its edge from the graph.
- **Calculate Result:**
  - After all pruning, the number of edges in the minimal required subtree is `remainingNodes - 1`.
  - The total travel distance is `2 * (remainingNodes - 1)`. Return `max(0, result)`.

# Solutions
### Java

```java
class Solution {
public
  int collectTheCoins(int[] coins, int[][] edges) {
    int n = coins.length;
    Set<Integer>[] g = new Set[n];
    Arrays.setAll(g, k->new HashSet<>());
    for (var e : edges) {
      int a = e[0], b = e[1];
      g[a].add(b);
      g[b].add(a);
    }
    Deque<Integer> q = new ArrayDeque<>();
    for (int i = 0; i < n; ++i) {
      if (coins[i] == 0 && g[i].size() == 1) {
        q.offer(i);
      }
    }
    while (!q.isEmpty()) {
      int i = q.poll();
      for (int j : g[i]) {
        g[j].remove(i);
        if (coins[j] == 0 && g[j].size() == 1) {
          q.offer(j);
        }
      }
      g[i].clear();
    }
    q.clear();
    for (int k = 0; k < 2; ++k) {
      for (int i = 0; i < n; ++i) {
        if (g[i].size() == 1) {
          q.offer(i);
        }
      }
      for (int i : q) {
        for (int j : g[i]) {
          g[j].remove(i);
        }
        g[i].clear();
      }
    }
    int ans = 0;
    for (var e : edges) {
      int a = e[0], b = e[1];
      if (g[a].size() > 0 && g[b].size() > 0) {
        ans += 2;
      }
    }
    return ans;
  }
}

```

### CPP

```cpp
class Solution {
public:
  int collectTheCoins(vector<int> &coins, vector<vector<int>> &edges) {
    int n = coins.size();
    unordered_set<int> g[n];
    for (auto &e : edges) {
      int a = e[0], b = e[1];
      g[a].insert(b);
      g[b].insert(a);
    }
    queue<int> q;
    for (int i = 0; i < n; ++i) {
      if (coins[i] == 0 && g[i].size() == 1) {
        q.push(i);
      }
    }
    while (!q.empty()) {
      int i = q.front();
      q.pop();
      for (int j : g[i]) {
        g[j].erase(i);
        if (coins[j] == 0 && g[j].size() == 1) {
          q.push(j);
        }
      }
      g[i].clear();
    }
    for (int k = 0; k < 2; ++k) {
      vector<int> q;
      for (int i = 0; i < n; ++i) {
        if (g[i].size() == 1) {
          q.push_back(i);
        }
      }
      for (int i : q) {
        for (int j : g[i]) {
          g[j].erase(i);
        }
        g[i].clear();
      }
    }
    int ans = 0;
    for (auto &e : edges) {
      int a = e[0], b = e[1];
      if (g[a].size() && g[b].size()) {
        ans += 2;
      }
    }
    return ans;
  }
};

```

### Python

```python
class Solution:
    def collectTheCoins(self, coins: List[int], edges: List[List[int]]) -> int: g = defaultdict(set) for a, b in edges: g[a]. add(b) g[b]. add(a) n = len(coins) q = deque(i for i in range(n) if len(g[i]) == 1 and coins[i] == 0) while q: i = q . popleft() for j in g[i]: g[j]. remove(i) if coins[j] == 0 and len(g[j]) == 1: q . append(j) g[i]. clear() for k in range(2): q = [i for i in range(n) if len(g[i]) == 1] for i in q: for j in g[i]: g[j]. remove(i) g[i]. clear() return sum(len(g[a]) > 0 and len(g[b]) > 0 for a, b in edges) * 2

```