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];
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;
        graph[i] = malloc(n * sizeof(int));
    }
    
    // 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) {
        // Free memory
        for (int i = 0; i < n; i++) {
            free(graph[i]);
        }
        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++;
        }
    }
    
    // Free memory
    for (int i = 0; i < n; i++) {
        free(graph[i]);
    }
    
    // We need to traverse each remaining edge twice (go and return)
    return 2 * edgeCount;
}

int main() {
    char line[100000];
    
    // Read coins - parse JSON array format [1,0,0,0,0,1]
    fgets(line, sizeof(line), stdin);
    int coinsSize = 0;
    char *ptr = line;
    
    // Find opening bracket
    while (*ptr && *ptr != '[') ptr++;
    if (*ptr == '[') ptr++;
    
    while (*ptr && *ptr != ']') {
        if (*ptr >= '0' && *ptr <= '9') {
            coins[coinsSize++] = *ptr - '0';
        }
        ptr++;
    }
    
    // Read edges - parse JSON 2D array format [[0,1],[1,2],[2,3],[3,4],[4,5]]
    fgets(line, sizeof(line), stdin);
    int edgesSize = 0;
    ptr = line;
    
    // Handle empty array case
    if (strstr(line, "[]") != NULL) {
        edgesSize = 0;
    } else {
        while (*ptr) {
            if (*ptr == '[') {
                // Check if this is the outer bracket
                char *nextPtr = ptr + 1;
                while (*nextPtr == ' ') nextPtr++;
                if (*nextPtr == '[') {
                    // This is the outer bracket, skip it
                    ptr++;
                    continue;
                }
                
                // This is an inner bracket, parse the pair
                ptr++;
                int a = 0, b = 0;
                
                // Read first number
                while (*ptr >= '0' && *ptr <= '9') {
                    a = a * 10 + (*ptr - '0');
                    ptr++;
                }
                
                // Skip comma and spaces
                while (*ptr && (*ptr < '0' || *ptr > '9')) {
                    if (*ptr == ']') break;
                    ptr++;
                }
                
                // Read second number if we haven't hit the closing bracket
                if (*ptr != ']') {
                    while (*ptr >= '0' && *ptr <= '9') {
                        b = b * 10 + (*ptr - '0');
                        ptr++;
                    }
                }
                
                edges[edgesSize][0] = a;
                edges[edgesSize][1] = b;
                edgesSize++;
                
                // Skip to next opening bracket or end
                while (*ptr && *ptr != '[' && *ptr != '\0') ptr++;
            } 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

Collect Coins in a Tree - Problem

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

DFS Traversal of All Nodes — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler