Find Number of Coins to Place in Tree Nodes

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.

• 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, say dfs(u, parent), will return the complete list of costs from the subtree rooted at u.
• Inside dfs(u, parent):
  • Initialize a list for u's subtree costs, my_costs, and add cost[u].
  • For each child v of u, recursively call dfs(v, u) to get the list of costs from v's subtree.
  • Merge the returned list from each child into my_costs.
  • The size of my_costs is now the size of u's subtree.
  • If the size is less than 3, coins[u] = 1.
  • Otherwise, sort my_costs and calculate the maximum product of three elements as described in the previous approach.
  • Store the result in coins[u].
  • Return the my_costs list to the caller (i.e., u's parent).

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;
    }
}