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.
This approach is a straightforward brute-force method. For every node in the tree, we explicitly determine its subtree and then calculate the required number of coins. This involves running a separate graph traversal for each of the n nodes.
Algorithm
- First, build an adjacency list representation for the undirected tree.
- To easily identify subtrees, convert this into a directed tree (parent-to-child relationships) rooted at 0. This can be done with a single Breadth-First or Depth-First Search starting from the root.
- Initialize a
coinsarray of sizen. - Iterate through each node
ifrom0ton-1. - For each node
i, perform a new traversal (e.g., DFS) starting fromion the directed tree to find all nodes in its subtree. - Collect the costs of all nodes in the subtree into a list.
- If the size of the subtree is less than 3, set
coins[i] = 1. - Otherwise, sort the list of costs.
- The maximum product can be either the product of the three largest numbers or the product of the two smallest (most negative) numbers and the single largest number.
- Calculate these two products:
product1 = cost_n-1 * cost_n-2 * cost_n-3andproduct2 = cost_0 * cost_1 * cost_n-1from the sorted list. - Set
coins[i]to the maximum of these two products. If the result is negative, set it to0.
Walkthrough
The core idea is to treat each node independently. We first process the edges to build a more convenient graph representation, like an adjacency list. Since the problem specifies a tree rooted at node 0, we can establish clear parent-child relationships. A good way to do this is to perform a single traversal (like BFS) from the root (node 0) to build a directed graph where edges only point from parent to child.
After setting up this directed tree structure, we iterate through every node i from 0 to n-1. For each i, we launch another traversal (like DFS) starting from i to visit all its descendants, thereby identifying its complete subtree. We gather the costs of all nodes in this subtree into a temporary list.
If the subtree size is less than 3, we simply place 1 coin. Otherwise, we sort the list of costs. The maximum product of three costs will come from either the three largest values or the two most negative values and the largest positive value. After sorting, these candidates are easy to find. We compute both potential products and take the larger one. If this maximum product is negative, we place 0 coins as per the problem rules.
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]); } List<Integer>[] children = new List[n]; for (int i = 0; i < n; i++) children[i] = new ArrayList<>(); Queue<Integer> q = new LinkedList<>(); q.offer(0); boolean[] visited = new boolean[n]; visited[0] = true; while (!q.isEmpty()) { int u = q.poll(); for (int v : adj[u]) { if (!visited[v]) { visited[v] = true; children[u].add(v); q.offer(v); } } } long[] coins = new long[n]; for (int i = 0; i < n; i++) { List<Long> subtreeCosts = new ArrayList<>(); getSubtreeCosts(i, children, cost, subtreeCosts); if (subtreeCosts.size() < 3) { coins[i] = 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[i] = Math.max(0L, Math.max(prod1, prod2)); } } return coins; } private void getSubtreeCosts(int u, List<Integer>[] children, int[] cost, List<Long> subtreeCosts) { subtreeCosts.add((long)cost[u]); for (int v : children[u]) { getSubtreeCosts(v, children, cost, subtreeCosts); } }}Complexity
Time
O(N^2 log N). Building the directed tree takes O(N). Then, for each of the N nodes, we traverse its subtree which can have up to N nodes, taking O(N) time. Sorting the collected costs takes O(N log N) in the worst case. This leads to a total complexity of O(N * (N + N log N)) = O(N^2 log N).
Space
O(N), where N is the number of nodes. This is for storing the adjacency list, the directed tree structure, and the list of subtree costs for each node (one at a time).
Trade-offs
Pros
Conceptually simple and easy to follow.
Cons
Extremely inefficient due to repeated computations.
The time complexity of O(N^2 log N) is too high for the given constraints and will likely result in a 'Time Limit Exceeded' error.
This is the most efficient approach, optimizing the DFS by avoiding passing large lists. The key observation is that for calculating the maximum product of three numbers, we only need the most extreme values: the three largest positive costs and the two smallest negative costs. Therefore, each node only needs to pass a small, constant-sized list of these candidate costs up to its parent.
Algorithm
- Build an adjacency list for the tree.
- Perform a post-order DFS. The DFS function for a node
uwill return a pair of values:(subtree_size, extreme_costs), whereextreme_costsis a small, sorted list of candidate costs. - Inside
dfs(u, parent):- Initialize
subtree_size = 1and a listcandidateswithcost[u]. - For each child
v, recursively calldfs(v, u)to get(child_size, child_extremes). - Add
child_sizetosubtree_sizeand mergechild_extremesintocandidates. - After visiting all children,
candidatesholdscost[u]and the extreme costs from all children subtrees. This list contains all values needed to find the overall extremes foru's subtree. - If
subtree_size < 3, setcoins[u] = 1. - Otherwise, sort
candidates. Calculate the max product using the 3 largest and 2 smallest values from this sorted list. Store the result incoins[u]. - Before returning, trim the
candidateslist. If its size is greater than 5, create a new list containing only its 2 smallest and 3 largest elements. - Return the pair
(subtree_size, trimmed_list).
- Initialize
Walkthrough
This method refines the post-order DFS strategy. Instead of returning the entire list of subtree costs, our recursive dfs function will return two things: the total size of the subtree, and a small, sorted list containing at most 5 key values from that subtree—the 2 smallest and 3 largest costs.
For a node u, the dfs function first calls itself on all children. It receives the subtree size and a small list of extreme costs from each child. It aggregates these sizes to compute its own subtree size. It also collects its own cost cost[u] and all the small lists from its children into a single temporary list, candidates.
This candidates list is guaranteed to contain the true 2 smallest and 3 largest costs for the entire subtree of u. We sort this small candidates list (its size is proportional to the degree of u, not its subtree size) and use it to calculate the coin value for u.
Finally, before returning to its parent, the function trims the candidates list down to the 2 smallest and 3 largest values. This ensures that only a constant amount of information is passed up at each step of the recursion, making the approach highly efficient in both time and memory.
class Solution { class Result { int size; List<Long> extremes; Result(int size, List<Long> extremes) { this.size = size; this.extremes = extremes; } } 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 Result dfs(int u, int p, List<Integer>[] adj, int[] cost, long[] coins) { int subtreeSize = 1; List<Long> candidates = new ArrayList<>(); candidates.add((long)cost[u]); for (int v : adj[u]) { if (v == p) continue; Result childResult = dfs(v, u, adj, cost, coins); subtreeSize += childResult.size; candidates.addAll(childResult.extremes); } if (subtreeSize < 3) { coins[u] = 1; } else { Collections.sort(candidates); int n_cand = candidates.size(); long prod1 = candidates.get(n_cand - 1) * candidates.get(n_cand - 2) * candidates.get(n_cand - 3); long prod2 = candidates.get(0) * candidates.get(1) * candidates.get(n_cand - 1); coins[u] = Math.max(0L, Math.max(prod1, prod2)); } Collections.sort(candidates); List<Long> nextExtremes; if (candidates.size() <= 5) { nextExtremes = candidates; } else { nextExtremes = new ArrayList<>(); int n_cand = candidates.size(); nextExtremes.add(candidates.get(0)); nextExtremes.add(candidates.get(1)); nextExtremes.add(candidates.get(n_cand - 3)); nextExtremes.add(candidates.get(n_cand - 2)); nextExtremes.add(candidates.get(n_cand - 1)); } return new Result(subtreeSize, nextExtremes); }}Complexity
Time
O(N log N). Each node is visited once. At a node `u`, the work is dominated by sorting the `candidates` list. The size of this list is at most `1 + 5 * (degree(u) - 1)`, so sorting takes `O(degree(u) * log(degree(u)))`. The total time is the sum over all nodes, `Σ O(deg(u)log(deg(u)))`, which is bounded by `O(N log N)` (the worst case being a star graph).
Space
O(N). The space is used for the adjacency list and the recursion stack. The depth of the recursion can be up to N, but the data stored per stack frame is constant-sized, leading to O(N) total space.
Trade-offs
Pros
Highly efficient time complexity, suitable for large inputs.
Optimal space complexity, as it avoids storing and passing large data structures.
Cons
The implementation is more complex than the previous approaches.
Requires careful handling of list manipulation and the data structure returned from the recursive calls.
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; }}class Solution {public: vector<long long> placedCoins(vector<vector<int>> &edges, vector<int> &cost) { int n = cost.size(); vector<long long> ans(n, 1); vector<int> g[n]; for (auto &e : edges) { int a = e[0], b = e[1]; g[a].push_back(b); g[b].push_back(a); } function<vector<int>(int, int)> dfs = [&](int a, int fa) -> vector<int> { vector<int> res = {cost[a]}; for (int b : g[a]) { if (b != fa) { auto t = dfs(b, a); res.insert(res.end(), t.begin(), t.end()); } } sort(res.begin(), res.end()); int m = res.size(); if (m >= 3) { long long x = 1LL * res[m - 1] * res[m - 2] * res[m - 3]; long long y = 1LL * res[0] * res[1] * res[m - 1]; ans[a] = max({0LL, x, y}); } if (m >= 5) { res = {res[0], res[1], res[m - 1], res[m - 2], res[m - 3]}; } return res; }; dfs(0, -1); return ans; }};class Solution: def placedCoins(self, edges: List[List[int]], cost: List[int]) -> List[int]: def dfs(a: int, fa: int) -> List[int]: res = [cost[a]] for b in g[a]: if b != fa: res . extend(dfs(b, a)) res . sort() if len(res) >= 3: ans[a] = max(res[- 3] * res[- 2] * res[- 1], res[0] * res[1] * res[- 1], 0) if len(res) > 5: res = res[: 2] + res[- 3:] return res n = len(cost) g = [[] for _ in range(n)] for a, b in edges: g[a]. append(b) g[b]. append(a) ans = [1] * n dfs(0, - 1) return ansVideo 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.