Skip to content

2973. Find Number of Coins to Place in Tree Nodes

SourceBiweekly Contest 120 Q4DifficultyHardRating2276

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 than 3, place 1 coin.
  • Otherwise, place an amount of coins equal to the maximum product of cost values assigned to 3 distinct nodes in the subtree of node i. If this product is negative, place 0 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

Thinking

A subtree with fewer than three nodes places \(1\); otherwise it places the maximum product of three costs, or \(0\) if that product is negative. The maximum is either the three largest values or the two smallest (negative) times the largest. A subtree may have \(n\) nodes, so the full list cannot be sent upward.

DFS returns a sorted list keeping the two smallest and three largest values. After merging children, trim, then compute the node’s answer.

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
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 ans

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

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

 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
func placedCoins(edges [][]int, cost []int) []int64 {
    n := len(cost)
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    ans := make([]int64, n)
    for i := range ans {
        ans[i] = int64(1)
    }
    var dfs func(a, fa int) []int
    dfs = func(a, fa int) []int {
        res := []int{cost[a]}
        for _, b := range g[a] {
            if b != fa {
                res = append(res, dfs(b, a)...)
            }
        }
        sort.Ints(res)
        m := len(res)
        if m >= 3 {
            x := res[m-1] * res[m-2] * res[m-3]
            y := res[0] * res[1] * res[m-1]
            ans[a] = max(0, int64(x), int64(y))
        }
        if m >= 5 {
            res = append(res[:2], res[m-3:]...)
        }
        return res
    }
    dfs(0, -1)
    return ans
}

 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
function placedCoins(edges: number[][], cost: number[]): number[] {
    const n = cost.length;
    const ans: number[] = Array(n).fill(1);
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const dfs = (a: number, fa: number): number[] => {
        const res: number[] = [cost[a]];
        for (const b of g[a]) {
            if (b !== fa) {
                res.push(...dfs(b, a));
            }
        }
        res.sort((a, b) => a - b);
        const m = res.length;
        if (m >= 3) {
            const x = res[m - 1] * res[m - 2] * res[m - 3];
            const y = res[0] * res[1] * res[m - 1];
            ans[a] = Math.max(0, x, y);
        }
        if (m > 5) {
            res.splice(2, m - 5);
        }
        return res;
    };
    dfs(0, -1);
    return ans;
}

Solution 2: Explicit Stack + Sorting

Thinking

A subtree with fewer than three nodes places \(1\) coin; otherwise it places the maximum product of three costs, or \(0\) when that product is negative. The maximum is either the three largest costs or the two smallest costs times the largest. A subtree may contain \(n\) nodes, so the full list cannot be sent upward, and each subtree keeps only the two smallest and three largest costs. Merging those lists by recursion along a chain uses a call depth equal to the node count and overflows Python once \(n\) reaches \(1000\), while \(n\) can be \(2 \times 10^4\). The stack stores \((node, parent, state)\): state \(0\) pushes the exit marker and then the children, and state \(1\) merges, sorts, and trims the child lists before writing the current answer.

There are two cases for the number of coins placed at node \(a\):

  • If the subtree of \(a\) has fewer than \(3\) nodes, place \(1\) coin.
  • If the subtree has at least \(3\) nodes, take \(3\) different nodes, place their maximum cost product at \(a\), and place \(0\) when that product is negative.

When every cost is positive, the product uses the \(3\) largest costs. When negatives appear, the two smallest costs times the largest cost must also be considered. Each subtree therefore keeps only its \(2\) smallest costs and its \(3\) largest costs.

Build the graph from \(edges\) and set every answer to \(1\). An explicit stack processes nodes in postorder. Each frame is \((a, fa, state)\). When \(state = 0\), push \((a, fa, 1)\) and then push every neighbor other than the parent with state \(0\). When \(state = 1\), merge \(cost[a]\) with the already trimmed lists of the children and sort the result into \(res\) of length \(m\):

  • If \(m \ge 3\), the number of coins at \(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 it stays \(1\).
  • If \(m > 5\), keep only the first \(2\) elements and the last \(3\) elements of \(res\) for the parent to merge.

The merged list is sorted again, so the order in which children are pushed does not change the answer.

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
23
24
25
26
27
28
29
30
31
class Solution:
    def placedCoins(self, edges: List[List[int]], cost: List[int]) -> List[int]:
        n = len(cost)
        g = [[] for _ in range(n)]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        ans = [1] * n
        sub = [None] * n
        stk = [(0, -1, 0)]
        while stk:
            a, fa, state = stk.pop()
            if state == 0:
                stk.append((a, fa, 1))
                for b in g[a]:
                    if b != fa:
                        stk.append((b, a, 0))
            else:
                res = [cost[a]]
                for b in g[a]:
                    if b != fa:
                        res.extend(sub[b])
                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:]
                sub[a] = res
        return ans

 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
43
44
45
46
47
48
49
50
class Solution {
    public long[] placedCoins(int[][] edges, int[] cost) {
        int n = cost.length;
        List<Integer>[] g = new List[n];
        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);
        }
        long[] ans = new long[n];
        Arrays.fill(ans, 1);
        List<Integer>[] sub = new List[n];
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {0, -1, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int a = cur[0], fa = cur[1], state = cur[2];
            if (state == 0) {
                stk.push(new int[] {a, fa, 1});
                for (int b : g[a]) {
                    if (b != fa) {
                        stk.push(new int[] {b, a, 0});
                    }
                }
            } else {
                List<Integer> res = new ArrayList<>();
                res.add(cost[a]);
                for (int b : g[a]) {
                    if (b != fa) {
                        res.addAll(sub[b]);
                    }
                }
                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 = new ArrayList<>(List.of(
                        res.get(0), res.get(1), res.get(m - 3), res.get(m - 2), res.get(m - 1)));
                }
                sub[a] = res;
            }
        }
        return ans;
    }
}

 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
43
44
45
46
class Solution {
public:
    vector<long long> placedCoins(vector<vector<int>>& edges, vector<int>& cost) {
        int n = cost.size();
        vector<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);
        }
        vector<long long> ans(n, 1);
        vector<vector<int>> sub(n);
        vector<array<int, 3>> stk{{0, -1, 0}};
        while (!stk.empty()) {
            auto [a, fa, state] = stk.back();
            stk.pop_back();
            if (state == 0) {
                stk.push_back({a, fa, 1});
                for (int b : g[a]) {
                    if (b != fa) {
                        stk.push_back({b, a, 0});
                    }
                }
            } else {
                vector<int> res = {cost[a]};
                for (int b : g[a]) {
                    if (b != fa) {
                        res.insert(res.end(), sub[b].begin(), sub[b].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 - 3], res[m - 2], res[m - 1]};
                }
                sub[a] = std::move(res);
            }
        }
        return ans;
    }
};

 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
43
44
45
46
47
48
func placedCoins(edges [][]int, cost []int) []int64 {
    n := len(cost)
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    ans := make([]int64, n)
    for i := range ans {
        ans[i] = 1
    }
    sub := make([][]int, n)
    type frame struct{ a, fa, state int }
    stk := []frame{{0, -1, 0}}
    for len(stk) > 0 {
        cur := stk[len(stk)-1]
        stk = stk[:len(stk)-1]
        a, fa, state := cur.a, cur.fa, cur.state
        if state == 0 {
            stk = append(stk, frame{a, fa, 1})
            for _, b := range g[a] {
                if b != fa {
                    stk = append(stk, frame{b, a, 0})
                }
            }
        } else {
            res := []int{cost[a]}
            for _, b := range g[a] {
                if b != fa {
                    res = append(res, sub[b]...)
                }
            }
            sort.Ints(res)
            m := len(res)
            if m >= 3 {
                x := int64(res[m-1]) * int64(res[m-2]) * int64(res[m-3])
                y := int64(res[0]) * int64(res[1]) * int64(res[m-1])
                ans[a] = max(x, y, int64(0))
            }
            if m > 5 {
                res = []int{res[0], res[1], res[m-3], res[m-2], res[m-1]}
            }
            sub[a] = res
        }
    }
    return ans
}

 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
function placedCoins(edges: number[][], cost: number[]): number[] {
    const n = cost.length;
    const ans: number[] = Array(n).fill(1);
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const sub: number[][] = Array(n);
    const stk: number[][] = [[0, -1, 0]];
    while (stk.length) {
        const [a, fa, state] = stk.pop()!;
        if (state === 0) {
            stk.push([a, fa, 1]);
            for (const b of g[a]) {
                if (b !== fa) {
                    stk.push([b, a, 0]);
                }
            }
        } else {
            const res: number[] = [cost[a]];
            for (const b of g[a]) {
                if (b !== fa) {
                    res.push(...sub[b]);
                }
            }
            res.sort((x, y) => x - y);
            const m = res.length;
            if (m >= 3) {
                const x = res[m - 1] * res[m - 2] * res[m - 3];
                const y = res[0] * res[1] * res[m - 1];
                ans[a] = Math.max(0, x, y);
            }
            if (m > 5) {
                sub[a] = [res[0], res[1], res[m - 3], res[m - 2], res[m - 1]];
            } else {
                sub[a] = res;
            }
        }
    }
    return ans;
}

Comments