LeetCode 2973 - Find Number of Coins to Place in Tree Nodes

Difficulty: 🔴 Hard
Topics: Dynamic Programming, Tree, Depth-First Search, Sorting, Heap (Priority Queue)

Solution

Problem Understanding

The problem gives us a tree with n nodes labeled from 0 to n - 1. The tree is undirected and rooted at node 0. The tree structure is described by the edges array, where each edge [a, b] connects nodes a and b. Each node also has an associated cost, given in the cost array.

Our task is to calculate the number of coins to place on each node according to specific rules:

If a node's subtree (the node itself plus all its descendants) contains fewer than 3 nodes, we place 1 coin on that node.
If a node's subtree contains 3 or more nodes, we need to select 3 distinct nodes in the subtree such that the product of their costs is maximized. If the maximum product is negative, we place 0 coins. Otherwise, the number of coins placed equals this maximum product.

The output should be an array coin of size n, where coin[i] is the number of coins placed at node i.

Constraints inform us that n can be up to 2 * 10^4, so any naive solution that examines all possible triplets in each subtree will be too slow. Additionally, costs may be negative, so we must carefully consider combinations that maximize the product, including cases where negative numbers can produce a positive product.

Important edge cases include:

Leaf nodes, where the subtree size is 1, automatically resulting in 1 coin.
Subtrees where the maximum product is negative, yielding 0 coins.
Subtrees with exactly 3 nodes, which is the minimal number required to calculate a product.

Approaches

Brute Force

A brute force approach would compute, for each node, the list of all nodes in its subtree and then iterate over all possible triplets of nodes to calculate the maximum product. This ensures correctness because it directly follows the problem's rules.

However, this is extremely inefficient. If a subtree has k nodes, the number of triplets is C(k, 3) = k*(k-1)*(k-2)/6. In the worst case, the root node has n nodes, producing O(n^3) operations, which is infeasible for n = 2 * 10^4.

Optimal Approach

The key insight is that we do not need all combinations of nodes to find the maximum product. For each node's subtree, only the three largest values and two smallest values matter, because the maximum product can come either from the top three largest costs or from a combination of the two most negative values with the largest positive value.

We can solve this problem efficiently using a post-order DFS traversal:

For each node, recursively gather the top three maximum costs and bottom two minimum costs from its children.
Combine them with the current node’s cost to find the maximum product in the subtree.
Count the size of the subtree to decide if we should place 1 coin (subtree size < 3) or the calculated product.

This reduces the complexity to O(n) because each node is visited once, and combining maximum/minimum values is a constant-time operation.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n^3)	O(n)	Iterates all triplets in subtree, too slow for n = 2e4
Optimal	O(n)	O(n)	Uses post-order DFS and tracks top 3 max / 2 min values for each subtree

Algorithm Walkthrough

Construct an adjacency list for the tree using the edges array.
Initialize a result array coin of size n.
Define a recursive DFS function that, for a node u and parent p, performs the following:
Initialize lists max_vals and min_vals to track the top 3 maximum costs and bottom 2 minimum costs in the subtree.
Initialize subtree_size to 1 (counting the current node itself).
For each child v of u (excluding the parent):
Recurse with DFS on v.
Merge v's top 3 max and bottom 2 min values into u’s max_vals and min_vals.
Increment subtree_size by v’s subtree size.
Append u’s own cost to both max_vals and min_vals.
Sort and trim max_vals to top 3, and min_vals to bottom 2.
If subtree_size < 3, set coin[u] = 1.
Otherwise, calculate all possible products using combinations of three numbers from max_vals and min_vals that could yield a maximum product.
Set coin[u] to the maximum product if it is non-negative, otherwise 0.
Call DFS starting from the root node 0.
Return the coin array.

Why it works: By propagating only the top three maximum costs and bottom two minimum costs from each child, we guarantee that all combinations that could yield the maximum product for a node are considered. The post-order traversal ensures that all children are processed before their parent, so the parent node always has the correct subtree information.

Python Solution

from typing import List
from itertools import combinations

class Solution:
    def placedCoins(self, edges: List[List[int]], cost: List[int]) -> List[int]:
        n = len(cost)
        adj = [[] for _ in range(n)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)

        coin = [0] * n

        def dfs(u: int, parent: int) -> (List[int], List[int], int):
            max_vals = []
            min_vals = []
            subtree_size = 1

            for v in adj[u]:
                if v == parent:
                    continue
                child_max, child_min, child_size = dfs(v, u)
                max_vals.extend(child_max)
                min_vals.extend(child_min)
                subtree_size += child_size

            max_vals.append(cost[u])
            min_vals.append(cost[u])

            max_vals.sort(reverse=True)
            max_vals = max_vals[:3]
            min_vals.sort()
            min_vals = min_vals[:2]

            if subtree_size < 3:
                coin[u] = 1
            else:
                candidates = set()
                if len(max_vals) >= 3:
                    candidates.add(max_vals[0] * max_vals[1] * max_vals[2])
                if len(min_vals) >= 2 and len(max_vals) >= 1:
                    candidates.add(min_vals[0] * min_vals[1] * max_vals[0])
                max_product = max(candidates) if candidates else 0
                coin[u] = max(max_product, 0)

            return max_vals, min_vals, subtree_size

        dfs(0, -1)
        return coin

Implementation Walkthrough: We build the adjacency list for efficient traversal. The DFS function returns three things: the top three maximum costs, the bottom two minimum costs, and the subtree size. At each node, we calculate candidates for maximum product and store the result in coin. Sorting and slicing ensures constant-time combination calculations.

Go Solution

package main

func placedCoins(edges [][]int, cost []int) []int64 {
    n := len(cost)
    adj := make([][]int, n)
    for _, e := range edges {
        u, v := e[0], e[1]
        adj[u] = append(adj[u], v)
        adj[v] = append(adj[v], u)
    }

    coin := make([]int64, n)

    var dfs func(u, parent int) ([]int, []int, int)
    dfs = func(u, parent int) ([]int, []int, int) {
        maxVals := []int{}
        minVals := []int{}
        subtreeSize := 1

        for _, v := range adj[u] {
            if v == parent {
                continue
            }
            childMax, childMin, childSize := dfs(v, u)
            maxVals = append(maxVals, childMax...)
            minVals = append(minVals, childMin...)
            subtreeSize += childSize
        }

        maxVals = append(maxVals, cost[u])
        minVals = append(minVals, cost[u])

        // sort descending for maxVals
        sort.Slice(maxVals, func(i, j int) bool { return maxVals[i] > maxVals[j] })
        if len(maxVals) > 3 {
            maxVals = maxVals[:3]
        }

        // sort ascending for minVals
        sort.Ints(minVals)
        if len(minVals) > 2 {
            minVals = minVals[:2]
        }

        if subtreeSize < 3 {
            coin[u] = 1
        } else {
            candidates := []int{}
            if len(maxVals) >= 3 {
                candidates = append(candidates, maxVals[0]*maxVals[1]*maxVals[2])
            }
            if len(minVals) >= 2 && len(maxVals) >= 1 {
                candidates = append(candidates, minVals[0]*minVals[1]*maxVals[0])
            }
            maxProduct