Find Number of Coins to Place in Tree Nodes - Problem

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

You are given an undirected tree with n nodes labeled from 0 to n - 1, and rooted at node 0. You are given 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 a 0-indexed integer array cost of length n, where cost[i] is the cost assigned to the ith node.

You need to place some coins on every node of the tree. The number of coins to be placed at node i can be calculated as:

If size of the subtree of node i is less than 3, place 1 coin.
Otherwise, place an amount of coins equal to the maximum product of cost values assigned to 3 distinct nodes in the subtree of node i.
If this product is negative, place 0 coins.

Return an array coin of size n such that coin[i] is the number of coins placed at node i.

Input & Output

Example 1 — Basic Tree

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

› Output: [30,30,1,1]

💡 Note: Node 0 subtree {0,1,2,3} has max product 3×2×5=30. Node 1 subtree {1,2,3} has max product 3×2×5=30. Nodes 2,3 have subtree size <3, so get 1 coin each.

Example 2 — Negative Values

$ Input: edges = [[0,1],[0,2]], cost = [-2,3,-1]

› Output: [6,1,1]

💡 Note: Node 0 subtree {0,1,2} has costs [-2,3,-1]. Max product is (-2)×3×(-1)=6. Nodes 1,2 have subtree size 1, so get 1 coin each.

Example 3 — Single Node

$ Input: edges = [], cost = [5]

› Output: [1]

💡 Note: Only one node with subtree size 1 < 3, so it gets 1 coin.

Constraints

1 ≤ n ≤ 2 × 10⁴
edges.length == n - 1
0 ≤ a_i, b_i < n
-10⁴ ≤ cost[i] ≤ 10⁴

Visualization

Tap to expand

Asked in

G Google 15 a Amazon 12 f Meta 8

The key insight is to use DFS to traverse the tree and collect subtree costs, then sort them to efficiently find the maximum product of 3 distinct values. For subtrees with <3 nodes, place 1 coin; otherwise find max product and place 0 if negative. Best approach uses DFS with sorted selection - Time: O(n² log n), Space: O(n).

Common Approaches

✓ Brute Force with All Combinations

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

Build the tree structure, then for each node collect all nodes in its subtree. If subtree has ≥3 nodes, try all combinations of 3 nodes to find maximum product. This approach is straightforward but inefficient for large trees.

DFS with Sorted Selection

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

Perform DFS from root to build tree structure. For each node, collect all costs in its subtree, sort them, and intelligently select 3 values that give maximum product. This avoids checking all combinations by using sorting to identify candidates.

Brute Force with All Combinations — Algorithm Steps

Step 1: Build adjacency list from edges and perform DFS to collect subtree nodes for each node
Step 2: For each node, if subtree size < 3, place 1 coin, otherwise try all C(n,3) combinations
Step 3: Find maximum product among all combinations, place 0 if negative, otherwise place the product

Visualization

Tap to expand

Step-by-Step Walkthrough

Build Tree

Create adjacency list from edges

For Each Node

Collect all nodes in its subtree using DFS

Try Combinations

Test all C(n,3) combinations to find maximum product

Code -

solution.c — C

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

static int adj[5000][5000];
static int adjSizes[5000];
static int subtrees[5000][5000];
static int subtreeSizes[5000];

void dfs(int node, int parent, int root) {
    subtrees[root][subtreeSizes[root]++] = node;
    
    for (int i = 0; i < adjSizes[node]; i++) {
        int child = adj[node][i];
        if (child != parent) {
            dfs(child, node, root);
        }
    }
}

void computeSubtrees(int node, int parent) {
    subtreeSizes[node] = 0;
    dfs(node, parent, node);
    
    for (int i = 0; i < adjSizes[node]; i++) {
        int child = adj[node][i];
        if (child != parent) {
            computeSubtrees(child, node);
        }
    }
}

int* solution(int** edges, int edgesSize, int* edgesColSize, int* cost, int costSize, int* returnSize) {
    int n = costSize;
    *returnSize = n;
    int* result = (int*)malloc(n * sizeof(int));
    
    if (n == 1) {
        result[0] = 1;
        return result;
    }
    
    // Initialize adjacency list
    for (int i = 0; i < n; i++) {
        adjSizes[i] = 0;
        subtreeSizes[i] = 0;
    }
    
    // Build adjacency list
    for (int i = 0; i < edgesSize; i++) {
        int a = edges[i][0], b = edges[i][1];
        adj[a][adjSizes[a]++] = b;
        adj[b][adjSizes[b]++] = a;
    }
    
    // Compute subtrees for each node when tree is rooted at 0
    computeSubtrees(0, -1);
    
    // Compute result for each node
    for (int i = 0; i < n; i++) {
        if (subtreeSizes[i] < 3) {
            result[i] = 1;
        } else {
            long long maxProduct = LLONG_MIN;
            // Try all combinations of 3 nodes
            for (int a = 0; a < subtreeSizes[i]; a++) {
                for (int b = a + 1; b < subtreeSizes[i]; b++) {
                    for (int c = b + 1; c < subtreeSizes[i]; c++) {
                        long long product = (long long)cost[subtrees[i][a]] * cost[subtrees[i][b]] * cost[subtrees[i][c]];
                        if (product > maxProduct) {
                            maxProduct = product;
                        }
                    }
                }
            }
            result[i] = maxProduct < 0 ? 0 : (int)maxProduct;
        }
    }
    
    return result;
}

long parseNumber(char* str, int* pos) {
    long num = 0;
    int sign = 1;
    if (str[*pos] == '-') {
        sign = -1;
        (*pos)++;
    }
    while (str[*pos] >= '0' && str[*pos] <= '9') {
        num = num * 10 + (str[*pos] - '0');
        (*pos)++;
    }
    return num * sign;
}

int main() {
    char line1[10000], line2[10000];
    fgets(line1, sizeof(line1), stdin);
    fgets(line2, sizeof(line2), stdin);
    
    // Parse edges
    int** edges = NULL;
    int edgesSize = 0;
    
    if (strncmp(line1, "[]", 2) != 0) {
        // Count edges by counting opening brackets after the first one
        char* ptr = line1 + 1; // skip first [
        while ((ptr = strchr(ptr, '[')) != NULL) {
            edgesSize++;
            ptr++;
        }
        
        if (edgesSize > 0) {
            edges = (int**)malloc(edgesSize * sizeof(int*));
            for (int i = 0; i < edgesSize; i++) {
                edges[i] = (int*)malloc(2 * sizeof(int));
            }
            
            int edgeIdx = 0;
            ptr = line1 + 1; // skip first [
            while ((ptr = strchr(ptr, '[')) != NULL) {
                ptr++; // skip [
                int pos = 0;
                char temp[50];
                int j = 0;
                while (*ptr != ']' && *ptr != '\0') {
                    temp[j++] = *ptr++;
                }
                temp[j] = '\0';
                
                pos = 0;
                edges[edgeIdx][0] = (int)parseNumber(temp, &pos);
                while (temp[pos] == ',' || temp[pos] == ' ') pos++;
                edges[edgeIdx][1] = (int)parseNumber(temp, &pos);
                edgeIdx++;
            }
        }
    }
    
    // Parse cost
    int costSize = 0;
    if (strncmp(line2, "[]", 2) != 0) {
        costSize = 1;
        char* ptr = line2;
        while ((ptr = strchr(ptr, ',')) != NULL) {
            costSize++;
            ptr++;
        }
    }
    
    int* cost = (int*)malloc(costSize * sizeof(int));
    if (costSize > 0) {
        char* ptr = line2 + 1; // skip [
        for (int i = 0; i < costSize; i++) {
            int pos = 0;
            char temp[50];
            int j = 0;
            while (*ptr != ',' && *ptr != ']' && *ptr != '\0' && *ptr != '\n') {
                temp[j++] = *ptr++;
            }
            temp[j] = '\0';
            pos = 0;
            cost[i] = (int)parseNumber(temp, &pos);
            if (*ptr == ',') ptr++;
        }
    }
    
    int returnSize;
    int* result = solution(edges, edgesSize, NULL, cost, costSize, &returnSize);
    
    printf("[");
    for (int i = 0; i < returnSize; i++) {
        printf("%d", result[i]);
        if (i < returnSize - 1) printf(",");
    }
    printf("]\n");
    
    free(result);
    if (edges) {
        for (int i = 0; i < edgesSize; i++) {
            free(edges[i]);
        }
        free(edges);
    }
    free(cost);
    
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n⁴)

For each of n nodes, collect subtree nodes O(n), then try all O(n³) combinations of 3 nodes

✓ Linear Growth

Space Complexity

O(n²)

Storing subtree nodes for each node requires O(n²) space in worst case

⚠ Quadratic Space

8.5K Views

Medium Frequency

~35 min Avg. Time

234 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

Find Number of Coins to Place in Tree Nodes - Problem

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

Brute Force with All Combinations — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler