Collect Coins in a Tree

Collect Coins in a Tree - 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.

Input & Output

Example 1 — Linear Tree with End Coins

$ Input: coins = [1,0,0,0,0,1], edges = [[0,1],[1,2],[2,3],[3,4],[4,5]]

› Output: 6

💡 Note: Start at node 2. Collect coin at 0 (distance 2). Move to 4, collect coin at 5 (distance 1). Return to 2. Total moves: 2 + 2 + 2 = 6

Example 2 — Tree with Adjacent Coins

$ Input: coins = [0,0,1,0,0,0], edges = [[0,1],[1,2],[1,3],[3,4],[4,5]]

› Output: 2

💡 Note: Start at node 1. Collect coin at node 2 (distance 1). No need to visit other branches. Total moves: 1 + 1 = 2

Example 3 — Single Node with Coin

$ Input: coins = [1], edges = []

› Output: 0

💡 Note: Only one node with a coin. Start there, collect it immediately. No movement needed.

Constraints

1 ≤ coins.length ≤ 3 × 10⁴
0 ≤ coins[i] ≤ 1
edges.length == coins.length - 1
edges[i].length == 2
0 ≤ ai, bi < coins.length
The input represents a valid tree

Visualization

Tap to expand

Asked in

G Google 15 f Meta 12 a Amazon 8 M Microsoft 6

The key insight is to recognize that you only need to visit nodes that either have coins or are necessary to reach nodes with coins. The optimal approach uses tree trimming to remove unnecessary leaf nodes without coins, then calculates the minimum traversal of the remaining tree. Time: O(n), Space: O(n)

Common Approaches

✓ DFS Traversal of All Nodes

⏱️ Time: O(n) Space: O(n)

Perform a complete DFS traversal visiting all nodes regardless of coin locations. Count every edge traversal including backtracking steps.

Tree Trimming with Topological Sort

⏱️ Time: O(n) Space: O(n)

First trim leaf nodes that don't have coins and aren't needed to reach coins. Then calculate minimum traversal on the reduced tree using the insight that we need to visit nodes within 2 steps of coins.

DFS Traversal of All Nodes — Algorithm Steps

Build adjacency list from edges
Start DFS from any node
Visit all neighbors recursively
Count each edge traversal twice (going and returning)

Visualization

Tap to expand

Step-by-Step Walkthrough

Build Graph

Create adjacency list from edges

DFS All

Visit every node in tree with DFS

Count Edges

Each edge traversed twice (go + return)

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MAX_N 30005
static int graph[MAX_N][MAX_N];
static int graphSize[MAX_N];
static int remaining[MAX_N];
static int coins[MAX_N];
static int edges[MAX_N][2];

int solution(int coinsSize, int edgesSize) {
    int n = coinsSize;
    if (n <= 1) return 0;
    
    // Initialize graph
    for (int i = 0; i < n; i++) {
        graphSize[i] = 0;
        remaining[i] = 1;
    }
    
    // Build adjacency list
    for (int i = 0; i < edgesSize; i++) {
        int a = edges[i][0], b = edges[i][1];
        graph[a][graphSize[a]++] = b;
        graph[b][graphSize[b]++] = a;
    }
    
    // Remove leaf nodes without coins (do this twice since we can collect coins at distance 2)
    for (int round = 0; round < 2; round++) {
        int toRemove[MAX_N];
        int removeCount = 0;
        
        for (int node = 0; node < n; node++) {
            if (!remaining[node]) continue;
            
            // Count neighbors that are still remaining
            int neighborCount = 0;
            for (int i = 0; i < graphSize[node]; i++) {
                if (remaining[graph[node][i]]) {
                    neighborCount++;
                }
            }
            // If it's a leaf (degree 1) and has no coin, mark for removal
            if (neighborCount <= 1 && coins[node] == 0) {
                toRemove[removeCount++] = node;
            }
        }
        
        for (int i = 0; i < removeCount; i++) {
            remaining[toRemove[i]] = 0;
        }
    }
    
    // Count remaining nodes
    int remainingCount = 0;
    for (int i = 0; i < n; i++) {
        if (remaining[i]) remainingCount++;
    }
    
    // If no nodes remain, all coins were on leaves that got removed
    if (remainingCount <= 1) {
        return 0;
    }
    
    // Count edges in remaining subgraph
    int edgeCount = 0;
    for (int i = 0; i < edgesSize; i++) {
        if (remaining[edges[i][0]] && remaining[edges[i][1]]) {
            edgeCount++;
        }
    }
    
    // We need to traverse each remaining edge twice (go and return)
    return 2 * edgeCount;
}

int main() {
    char line[100000];
    
    // Read coins
    fgets(line, sizeof(line), stdin);
    int coinsSize = 0;
    char *ptr = line;
    while (*ptr) {
        if (*ptr >= '0' && *ptr <= '9') {
            coins[coinsSize++] = *ptr - '0';
        }
        ptr++;
    }
    
    // Read edges
    fgets(line, sizeof(line), stdin);
    int edgesSize = 0;
    ptr = line;
    while (*ptr) {
        if (*ptr == '[') {
            ptr++;
            int a = 0, b = 0;
            // Read first number
            while (*ptr >= '0' && *ptr <= '9') {
                a = a * 10 + (*ptr - '0');
                ptr++;
            }
            // Skip comma and space
            while (*ptr && (*ptr < '0' || *ptr > '9')) {
                ptr++;
            }
            // Read second number
            while (*ptr >= '0' && *ptr <= '9') {
                b = b * 10 + (*ptr - '0');
                ptr++;
            }
            edges[edgesSize][0] = a;
            edges[edgesSize][1] = b;
            edgesSize++;
        } else {
            ptr++;
        }
    }
    
    int result = solution(coinsSize, edgesSize);
    printf("%d\n", result);
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n)

DFS visits each node once and traverses each edge twice

✓ Linear Growth

Space Complexity

O(n)

Recursion stack depth up to n for skewed trees plus adjacency list

⚡ Linearithmic Space

23.0K Views

Medium Frequency

~35 min Avg. Time

890 Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

DFS Traversal of All Nodes — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler