Find Number of Coins to Place in Tree Nodes
HardPrompt
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
iis less than3, place1coin. - Otherwise, place an amount of coins equal to the maximum product of cost values assigned to
3distinct nodes in the subtree of nodei. If this product is negative, place0coins.
Return an array coin of size n such that coin[i] is the number of coins placed at node i.
Example 1:
Input: edges = [[0,1],[0,2],[0,3],[0,4],[0,5]], cost = [1,2,3,4,5,6]
Output: [120,1,1,1,1,1]
Explanation: For node 0 place 6 * 5 * 4 = 120 coins. All other nodes are leaves with subtree of size 1, place 1 coin on each of them.Example 2:
Input: edges = [[0,1],[0,2],[1,3],[1,4],[1,5],[2,6],[2,7],[2,8]], cost = [1,4,2,3,5,7,8,-4,2]
Output: [280,140,32,1,1,1,1,1,1]
Explanation: The coins placed on each node are:
- Place 8 * 7 * 5 = 280 coins on node 0.
- Place 7 * 5 * 4 = 140 coins on node 1.
- Place 8 * 2 * 2 = 32 coins on node 2.
- All other nodes are leaves with subtree of size 1, place 1 coin on each of them.Example 3:
Input: edges = [[0,1],[0,2]], cost = [1,2,-2]
Output: [0,1,1]
Explanation: Node 1 and 2 are leaves with subtree of size 1, place 1 coin on each of them. For node 0 the only possible product of cost is 2 * 1 * -2 = -4. Hence place 0 coins on node 0.
Constraints:
2 <= n <= 2 * 104edges.length == n - 1edges[i].length == 20 <= ai, bi < ncost.length == n1 <= |cost[i]| <= 104- The input is generated such that
edgesrepresents a valid tree.
Approaches
3 approaches with complexity analysis and trade-offs.
To improve upon the brute-force method, we can avoid recomputing subtree information by using a single post-order traversal (DFS). In this approach, each node processes its children first, receives the complete list of costs from their subtrees, aggregates them, computes its own coin value, and then passes the combined list up to its parent.
Algorithm
- Build an adjacency list for the tree.
- Perform a single post-order traversal (DFS) starting from the root (node 0).
- The DFS function for a node
u, saydfs(u, parent), will return the complete list of costs from the subtree rooted atu. - Inside
dfs(u, parent):- Initialize a list for
u's subtree costs,my_costs, and addcost[u]. - For each child
vofu, recursively calldfs(v, u)to get the list of costs fromv's subtree. - Merge the returned list from each child into
my_costs. - The size of
my_costsis now the size ofu's subtree. - If the size is less than 3,
coins[u] = 1. - Otherwise, sort
my_costsand calculate the maximum product of three elements as described in the previous approach. - Store the result in
coins[u]. - Return the
my_costslist to the caller (i.e.,u's parent).
- Initialize a list for
Walkthrough
This approach leverages the property of post-order traversal where children are visited before their parent. This allows information to flow up the tree. We define a recursive DFS function that, for a given node u, returns a list containing all the costs in the subtree of u.
When the DFS function is called on a node u, it first makes recursive calls on all its children. Each child returns a list of costs from its respective subtree. Node u then collects all these lists, adds its own cost cost[u], and merges them into a single comprehensive list. This list now represents all costs in u's subtree.
With this complete list, u can determine its subtree size and calculate the number of coins to be placed on it. If the size is less than 3, it's 1 coin. Otherwise, it sorts the list and finds the maximum product of three costs. Finally, the function returns the complete list of subtree costs to its parent, continuing the process up the tree.
While this avoids redundant traversals, it can be inefficient as the lists of costs passed up the tree can become very large, especially for nodes near the root.
class Solution { public long[] getCoinAmount(int[][] edges, int[] cost) { int n = cost.length; List<Integer>[] adj = new List[n]; for (int i = 0; i < n; i++) adj[i] = new ArrayList<>(); for (int[] edge : edges) { adj[edge[0]].add(edge[1]); adj[edge[1]].add(edge[0]); } long[] coins = new long[n]; dfs(0, -1, adj, cost, coins); return coins; } private List<Long> dfs(int u, int p, List<Integer>[] adj, int[] cost, long[] coins) { List<Long> subtreeCosts = new ArrayList<>(); subtreeCosts.add((long)cost[u]); for (int v : adj[u]) { if (v == p) continue; subtreeCosts.addAll(dfs(v, u, adj, cost, coins)); } if (subtreeCosts.size() < 3) { coins[u] = 1; } else { Collections.sort(subtreeCosts); int sz = subtreeCosts.size(); long prod1 = subtreeCosts.get(sz - 1) * subtreeCosts.get(sz - 2) * subtreeCosts.get(sz - 3); long prod2 = subtreeCosts.get(0) * subtreeCosts.get(1) * subtreeCosts.get(sz - 1); coins[u] = Math.max(0L, Math.max(prod1, prod2)); } return subtreeCosts; }}Complexity
Time
O(N^2) in the worst case. At each node `u`, we merge lists from its children. The total size of lists merged at `u` is proportional to its subtree size. The sum of all subtree sizes across all nodes in a line graph is `O(N^2)`. Sorting at each step would add a log factor, but the merging cost dominates.
Space
O(N^2) in the worst case. For a skewed tree (like a line), the recursion depth is N, and the lists passed up the stack grow in size from 1 to N. The total space on the call stack would be the sum of `1 + 2 + ... + N`, which is O(N^2).
Trade-offs
Pros
More efficient than the brute-force approach as it traverses the tree only once.
Eliminates redundant calculations for subtrees.
Cons
The time and space complexity are still high for large N.
Passing large lists through the recursion stack can be very memory-intensive, potentially leading to stack overflow or out-of-memory errors for deep or large-degree trees.
Solutions
Solution
class Solution {private int[] cost;private List<Integer>[] g;private long[] ans;public long[] placedCoins(int[][] edges, int[] cost) { int n = cost.length; this.cost = cost; ans = new long[n]; g = new List[n]; Arrays.fill(ans, 1); Arrays.setAll(g, i->new ArrayList<>()); for (int[] e : edges) { int a = e[0], b = e[1]; g[a].add(b); g[b].add(a); } dfs(0, -1); return ans; }private List<Integer> dfs(int a, int fa) { List<Integer> res = new ArrayList<>(); res.add(cost[a]); for (int b : g[a]) { if (b != fa) { res.addAll(dfs(b, a)); } } Collections.sort(res); int m = res.size(); if (m >= 3) { long x = (long)res.get(m - 1) * res.get(m - 2) * res.get(m - 3); long y = (long)res.get(0) * res.get(1) * res.get(m - 1); ans[a] = Math.max(0, Math.max(x, y)); } if (m >= 5) { res = List.of(res.get(0), res.get(1), res.get(m - 3), res.get(m - 2), res.get(m - 1)); } return res; }}Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.