Skip to content

Commit b2f4f91

Browse files
authored
Create 2603-collect-coins-in-a-tree.js
1 parent 6d7b138 commit b2f4f91

File tree

1 file changed

+35
-0
lines changed

1 file changed

+35
-0
lines changed

2603-collect-coins-in-a-tree.js

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
let a, cum, res, v, sum, g;
2+
const initializeGraph = (n) => { let g = []; for (let i = 0; i < n; i++) { g.push([]); } return g; };
3+
const packUG = (g, edges) => { for (const [u, v] of edges) { g[u].push(v); g[v].push(u); } };
4+
const sm = (a) => a.reduce(((x, y) => x + y), 0);
5+
const tree_dp = (cur, par) => {
6+
v[cur] = a[cur];
7+
for (const child of g[cur]) {
8+
if (child != par) {
9+
v[cur] += tree_dp(child, cur);
10+
}
11+
}
12+
if (cur != par) {
13+
let x = v[cur] + cum[par] - a[cur];
14+
let y = (sum - v[cur]) + cum[cur] - a[par];
15+
if (x != sum && y != sum) res += 2;
16+
}
17+
return v[cur];
18+
};
19+
/**
20+
* @param {number[]} coins
21+
* @param {number[][]} edges
22+
* @return {number}
23+
*/
24+
const collectTheCoins = function(coins, edges) {
25+
let n = coins.length;
26+
g = initializeGraph(n), a = coins, res = 0, cum = Array(n), v = Array(n), sum = sm(a);
27+
packUG(g, edges);
28+
for (let i = 0; i < n; i++) {
29+
cum[i] = a[i];
30+
for (const child of g[i]) cum[i] += a[child];
31+
}
32+
tree_dp(0, 0);
33+
return res;
34+
35+
};

0 commit comments

Comments
 (0)