2920. Maximum Points After Collecting Coins From All Nodes

Description

There exists an undirected tree rooted at node 0 with n nodes labeled from 0 to n - 1. You are given a 2D integer array edges of length n - 1, where edges[i] = [a_i, b_i] indicates that there is an edge between nodes a_i and b_i in the tree. You are also given a 0-indexed array coins of size n where coins[i] indicates the number of coins in the vertex i, and an integer k.

Starting from the root, you have to collect all the coins such that the coins at a node can only be collected if the coins of its ancestors have been already collected.

Coins at node_i can be collected in one of the following ways:

Collect all the coins, but you will get coins[i] - k points. If coins[i] - k is negative then you will lose abs(coins[i] - k) points.
Collect all the coins, but you will get floor(coins[i] / 2) points. If this way is used, then for all the node_j present in the subtree of node_i, coins[j] will get reduced to floor(coins[j] / 2).

Return the maximum points you can get after collecting the coins from all the tree nodes.

Example 1:

Input: edges = [[0,1],[1,2],[2,3]], coins = [10,10,3,3], k = 5
Output: 11                        
Explanation: 
Collect all the coins from node 0 using the first way. Total points = 10 - 5 = 5.
Collect all the coins from node 1 using the first way. Total points = 5 + (10 - 5) = 10.
Collect all the coins from node 2 using the second way so coins left at node 3 will be floor(3 / 2) = 1. Total points = 10 + floor(3 / 2) = 11.
Collect all the coins from node 3 using the second way. Total points = 11 + floor(1 / 2) = 11.
It can be shown that the maximum points we can get after collecting coins from all the nodes is 11.

Example 2:

Input: edges = [[0,1],[0,2]], coins = [8,4,4], k = 0
Output: 16
Explanation: 
Coins will be collected from all the nodes using the first way. Therefore, total points = (8 - 0) + (4 - 0) + (4 - 0) = 16.

Constraints:

n == coins.length
2 <= n <= 10⁵
0 <= coins[i] <= 10⁴
edges.length == n - 1
0 <= edges[i][0], edges[i][1] < n
0 <= k <= 10⁴

Solutions

Solution 1: Memoization Search

First, we construct a graph $g$ based on the edges given in the problem, where $g[i]$ represents all adjacent nodes of node $i$. Then we can use the method of memoization search to solve this problem.

We design a function $dfs(i, fa, j)$, which represents that the current node is $i$, the parent node is $fa$, the number of gold coins of the current node needs to be shifted to the right by $j$ bits, and the maximum score that can be obtained.

The execution process of the function $dfs(i, fa, j)$ is as follows:

If we use the first method to collect the gold coins of the current node, then the score of the current node is $(coins[i] » j) - k$. Then we traverse all the adjacent nodes $c$ of the current node. If $c$ is not equal to $fa$, then we add the result of $dfs(c, i, j)$ to the score of the current node.

If we use the second method to collect the gold coins of the current node, then the score of the current node is $coins[i] » (j + 1)$. Then we traverse all the adjacent nodes $c$ of the current node. If $c$ is not equal to $fa$, then we add the result of $dfs(c, i, j + 1)$ to the score of the current node. Note that since the maximum value of $coins[i]$ given in the problem is $10^4$, we can only shift to the right by at most $14$ bits, so that the value of $coins[i] » (j + 1)$ is $0$.

Finally, we return the maximum score that can be obtained by using the two methods at the current node.

In order to avoid repeated calculations, we use the method of memoization search and store the result of $dfs(i, fa, j)$ in $f[i][j]$, where $f[i][j]$ represents that the current node is $i$, the parent node is $fa$, the number of gold coins of the current node needs to be shifted to the right by $j$ bits, and the maximum score that can be obtained.

The time complexity is $O(n \times \log M)$, and the space complexity is $O(n \times \log M)$. Where $M$ represents the maximum value of $coins[i]$.

Solution 2: Tree DP + Explicit Stack

Thinking

At each node we may subtract $k$ or right-shift the coins once more. The shift applies to the unprocessed subtree, and coins are at most $10^4$, so about $14$ shifts wipe them out. With $n \le 10^5$, recursing from the root at the same shift is too deep: a chain makes the call depth $n$.

The best score of node $i$ after $j$ shifts depends only on the children’s scores after $j$ shifts and after $j + 1$ shifts. One choice takes $(coins[i] \gg j) - k$ and keeps shift $j$ below; the other takes $coins[i] \gg (j + 1)$ and uses shift $j + 1$, stopping once $j$ reaches $14$.

An explicit stack of $(node, parent, state)$ runs the postorder. After the children are left, we fill this node’s table for shifts $0$ through $14$. The answer is the root’s score at shift $0$.

First, we construct a graph $g$ based on the edges given in the problem, where $g[i]$ represents all adjacent nodes of node $i$.

$f[i][j]$ is the maximum score of the subtree at node $i$ after the coins have been shifted right by $j$ bits. An explicit stack walks from the root in postorder and fills a node only after its children are filled. For each shift $j$ with $0 \le j \le 14$:

If we collect the current node’s coins by the first method, its score is $(coins[i] » j) - k$, plus $f[c][j]$ for every child $c$.

If we collect them by the second method, its score is $coins[i] » (j + 1)$. When $j < 14$, we also add $f[c][j + 1]$ for every child. Coins are at most $10^4$, so after $14$ shifts the remainder is $0$.

$f[i][j]$ is the larger of the two methods. The answer is $f[0][0]$.

The time complexity is $O(n \times \log M)$, and the space complexity is $O(n \times \log M)$. Where $M$ represents the maximum value of $coins[i]$.

class Solution {
    private int k;
    private int[] coins;
    private Integer[][] f;
    private List<Integer>[] g;

    public int maximumPoints(int[][] edges, int[] coins, int k) {
        this.k = k;
        this.coins = coins;
        int n = coins.length;
        f = new Integer[n][15];
        g = new List[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        return dfs(0, -1, 0);
    }

    private int dfs(int i, int fa, int j) {
        if (f[i][j] != null) {
            return f[i][j];
        }
        int a = (coins[i] >> j) - k;
        int b = coins[i] >> (j + 1);
        for (int c : g[i]) {
            if (c != fa) {
                a += dfs(c, i, j);
                if (j < 14) {
                    b += dfs(c, i, j + 1);
                }
            }
        }
        return f[i][j] = Math.max(a, b);
    }
}


// Solution 2
class Solution {
    public int maximumPoints(int[][] edges, int[] coins, int k) {
        int n = coins.length;
        List<Integer>[] g = new List[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        int[][] f = new int[n][15];
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {0, -1, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int i = cur[0], fa = cur[1], state = cur[2];
            if (state == 0) {
                stk.push(new int[] {i, fa, 1});
                for (int c : g[i]) {
                    if (c != fa) {
                        stk.push(new int[] {c, i, 0});
                    }
                }
            } else {
                for (int j = 0; j < 15; ++j) {
                    int a = (coins[i] >> j) - k;
                    int b = coins[i] >> (j + 1);
                    for (int c : g[i]) {
                        if (c != fa) {
                            a += f[c][j];
                            if (j < 14) {
                                b += f[c][j + 1];
                            }
                        }
                    }
                    f[i][j] = Math.max(a, b);
                }
            }
        }
        return f[0][0];
    }
}

class Solution {
public:
    int maximumPoints(vector<vector<int>>& edges, vector<int>& coins, int k) {
        int n = coins.size();
        int f[n][15];
        memset(f, -1, sizeof(f));
        vector<int> g[n];
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].emplace_back(b);
            g[b].emplace_back(a);
        }
        function<int(int, int, int)> dfs = [&](int i, int fa, int j) {
            if (f[i][j] != -1) {
                return f[i][j];
            }
            int a = (coins[i] >> j) - k;
            int b = coins[i] >> (j + 1);
            for (int c : g[i]) {
                if (c != fa) {
                    a += dfs(c, i, j);
                    if (j < 14) {
                        b += dfs(c, i, j + 1);
                    }
                }
            }
            return f[i][j] = max(a, b);
        };
        return dfs(0, -1, 0);
    }
};


// Solution 2
class Solution {
public:
    int maximumPoints(vector<vector<int>>& edges, vector<int>& coins, int k) {
        int n = coins.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<array<int, 15>> f(n);
        vector<array<int, 3>> stk{{0, -1, 0}};
        while (!stk.empty()) {
            auto [i, fa, state] = stk.back();
            stk.pop_back();
            if (state == 0) {
                stk.push_back({i, fa, 1});
                for (int c : g[i]) {
                    if (c != fa) {
                        stk.push_back({c, i, 0});
                    }
                }
            } else {
                for (int j = 0; j < 15; ++j) {
                    int a = (coins[i] >> j) - k;
                    int b = coins[i] >> (j + 1);
                    for (int c : g[i]) {
                        if (c != fa) {
                            a += f[c][j];
                            if (j < 14) {
                                b += f[c][j + 1];
                            }
                        }
                    }
                    f[i][j] = max(a, b);
                }
            }
        }
        return f[0][0];
    }
};

class Solution:
    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:
        @cache
        def dfs(i: int, fa: int, j: int) -> int:
            a = (coins[i] >> j) - k
            b = coins[i] >> (j + 1)
            for c in g[i]:
                if c != fa:
                    a += dfs(c, i, j)
                    if j < 14:
                        b += dfs(c, i, j + 1)
            return max(a, b)

        n = len(coins)
        g = [[] for _ in range(n)]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        ans = dfs(0, -1, 0)
        dfs.cache_clear()
        return ans



# Solution 2
class Solution:
    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:
        n = len(coins)
        g = [[] for _ in range(n)]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        f = [[0] * 15 for _ in range(n)]
        stk = [(0, -1, 0)]
        while stk:
            i, fa, state = stk.pop()
            if state == 0:
                stk.append((i, fa, 1))
                for c in g[i]:
                    if c != fa:
                        stk.append((c, i, 0))
            else:
                for j in range(15):
                    a = (coins[i] >> j) - k
                    b = coins[i] >> (j + 1)
                    for c in g[i]:
                        if c != fa:
                            a += f[c][j]
                            if j < 14:
                                b += f[c][j + 1]
                    f[i][j] = max(a, b)
        return f[0][0]

func maximumPoints(edges [][]int, coins []int, k int) int {
	n := len(coins)
	f := make([][]int, n)
	for i := range f {
		f[i] = make([]int, 15)
		for j := range f[i] {
			f[i][j] = -1
		}
	}
	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)
	}
	var dfs func(int, int, int) int
	dfs = func(i, fa, j int) int {
		if f[i][j] != -1 {
			return f[i][j]
		}
		a := (coins[i] >> j) - k
		b := coins[i] >> (j + 1)
		for _, c := range g[i] {
			if c != fa {
				a += dfs(c, i, j)
				if j < 14 {
					b += dfs(c, i, j+1)
				}
			}
		}
		f[i][j] = max(a, b)
		return f[i][j]
	}
	return dfs(0, -1, 0)
}


// Solution 2
func maximumPoints(edges [][]int, coins []int, k int) int {
	n := len(coins)
	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)
	}
	f := make([][15]int, n)
	stk := [][3]int{{0, -1, 0}}
	for len(stk) > 0 {
		cur := stk[len(stk)-1]
		stk = stk[:len(stk)-1]
		i, fa, state := cur[0], cur[1], cur[2]
		if state == 0 {
			stk = append(stk, [3]int{i, fa, 1})
			for _, c := range g[i] {
				if c != fa {
					stk = append(stk, [3]int{c, i, 0})
				}
			}
		} else {
			for j := 0; j < 15; j++ {
				a := (coins[i] >> j) - k
				b := coins[i] >> (j + 1)
				for _, c := range g[i] {
					if c != fa {
						a += f[c][j]
						if j < 14 {
							b += f[c][j+1]
						}
					}
				}
				f[i][j] = max(a, b)
			}
		}
	}
	return f[0][0]
}

function maximumPoints(edges: number[][], coins: number[], k: number): number {
    const n = coins.length;
    const f: number[][] = Array.from({ length: n }, () => Array(15).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 = (i: number, fa: number, j: number): number => {
        if (f[i][j] !== -1) {
            return f[i][j];
        }
        let a = (coins[i] >> j) - k;
        let b = coins[i] >> (j + 1);
        for (const c of g[i]) {
            if (c !== fa) {
                a += dfs(c, i, j);
                if (j < 14) {
                    b += dfs(c, i, j + 1);
                }
            }
        }
        return (f[i][j] = Math.max(a, b));
    };
    return dfs(0, -1, 0);
}



// Solution 2
function maximumPoints(edges: number[][], coins: number[], k: number): number {
    const n = coins.length;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const f: number[][] = Array.from({ length: n }, () => Array(15).fill(0));
    const stk: [number, number, number][] = [[0, -1, 0]];
    while (stk.length) {
        const [i, fa, state] = stk.pop()!;
        if (state === 0) {
            stk.push([i, fa, 1]);
            for (const c of g[i]) {
                if (c !== fa) {
                    stk.push([c, i, 0]);
                }
            }
        } else {
            for (let j = 0; j < 15; ++j) {
                let a = (coins[i] >> j) - k;
                let b = coins[i] >> (j + 1);
                for (const c of g[i]) {
                    if (c !== fa) {
                        a += f[c][j];
                        if (j < 14) {
                            b += f[c][j + 1];
                        }
                    }
                }
                f[i][j] = Math.max(a, b);
            }
        }
    }
    return f[0][0];
}

2920 - Maximum Points After Collecting Coins From All Nodes

2920. Maximum Points After Collecting Coins From All Nodes

Description

Solutions

Solution 2: Tree DP + Explicit Stack

All Problems

All Solutions

Welcome to Subscribe On Youtube

2920. Maximum Points After Collecting Coins From All Nodes

Description

Solutions

Solution 2: Tree DP + Explicit Stack

All Problems

All Solutions