2973. Find Number of Coins to Place in Tree Nodes
Description
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 than3
, place1
coin. - Otherwise, place an amount of coins equal to the maximum product of cost values assigned to
3
distinct nodes in the subtree of nodei
. If this product is negative, place0
coins.
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 * 104
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
cost.length == n
1 <= |cost[i]| <= 104
- The input is generated such that
edges
represents a valid tree.
Solutions
Solution 1: DFS + Sorting
According to the problem description, there are two situations for the number of coins placed at each node $a$:
- If the number of nodes in the subtree corresponding to node $a$ is less than $3$, then place $1$ coin;
- If the number of nodes in the subtree corresponding to node $a$ is greater than or equal to $3$, then we need to take out $3$ different nodes from the subtree, calculate the maximum value of their cost product, and then place the corresponding number of coins at node $a$. If the maximum product is negative, place $0$ coins.
The first situation is relatively simple, we just need to count the number of nodes in the subtree of each node during the traversal.
For the second situation, if all costs are positive, we should take the $3$ nodes with the largest costs; if there are negative costs, we should take the $2$ nodes with the smallest costs and the $1$ node with the largest cost. Therefore, we need to maintain the smallest $2$ costs and the largest $3$ costs in each subtree.
We first construct the adjacency list $g$ based on the given two-dimensional array $edges$, where $g[a]$ represents all neighbor nodes of node $a$.
Next, we design a function $dfs(a, fa)$, which returns an array $res$, which stores the smallest $2$ costs and the largest $3$ costs in the subtree of node $a$ (may not be $5$).
In the function $dfs(a, fa)$, we add the cost $cost[a]$ of node $a$ to the array $res$, and then traverse all neighbor nodes $b$ of node $a$. If $b$ is not the parent node $fa$ of node $a$, then we add the result of $dfs(b, a)$ to the array $res$.
Then, we sort the array $res$, and then calculate the number of coins placed at node $a$ based on the length $m$ of the array $res$, and update $ans[a]$:
- If $m \ge 3$, then the number of coins placed at node $a$ is $\max(0, res[m - 1] \times res[m - 2] \times res[m - 3], res[0] \times res[1] \times res[m - 1])$, otherwise the number of coins placed at node $a$ is $1$;
- If $m > 5$, then we only need to keep the first $2$ elements and the last $3$ elements of the array $res$.
Finally, we call the function $dfs(0, -1)$, and return the answer array $ans$.
The time complexity is $O(n \times \log n)$, and the space complexity is $O(n)$. Where $n$ is the number of nodes.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
|