Maximum Points After Collecting Coins From All Nodes

Hard

#2601Time: O(2^N) in the worst case. At each node, the recursion branches into two possibilities. For a tree of depth `N`, this can lead to an exponential number of calls. Even with capping the number of halvings, the number of calls is very large without memoization.Space: O(N) for the recursion stack in the worst case (for a skewed tree) and for storing the adjacency list.1 company

Patterns

Dynamic Programming Bit Manipulation Memoization

Algorithms

Depth-First Search

Data structures

Array Tree

Companies

DE Shaw

Prompt

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⁴

Approaches

2 approaches with complexity analysis and trade-offs.

This approach uses a simple recursive function to explore all possible choices at each node. For every node, we recursively calculate the maximum points obtainable from its subtree for both possible actions:

Take coins[i] - k points.
Take floor(coins[i] / 2) points and halve the coins in the entire subtree.

The function will take the current node, its parent, and the number of times its coin value has been halved by its ancestors as parameters. It will make two recursive calls for each child, corresponding to the two choices at the current node, and sum up the results. The maximum of the two choices is returned.

Algorithm

Build an adjacency list for the tree from the edges array.
Define a recursive function solve(u, parent, h) which computes the maximum points from the subtree at u, given h prior halvings.
Inside solve(u, parent, h):
- Calculate the current coin value at u: current_coins = coins[u] >> h.
- Calculate the score for Option 1: (current_coins - k) + sum(solve(v, u, h)) for all children v.
- Calculate the score for Option 2: (current_coins / 2) + sum(solve(v, u, h + 1)) for all children v.
- Return the maximum of the two scores.
Start the recursion with solve(0, -1, 0).

Walkthrough

The core of this approach is a recursive function, let's call it solve(u, parent, h). This function calculates the maximum points we can get from the subtree rooted at node u, given that the coin values in this subtree have already been halved h times due to decisions made at u's ancestors.

The recursion works as follows:

First, we need to represent the tree. An adjacency list is a suitable choice. We build it from the input edges.
The recursive function solve(u, parent, h):
- Calculates the current coin value for node u: current_coins = coins[u] >> h. Note that >> h is equivalent to floor(c / 2^h).
- Option 1: We choose to get current_coins - k points from node u. The number of halvings for its children's subtrees remains h. We recursively call solve(v, u, h) for each child v of u and sum up the results.
- Option 2: We choose to get floor(current_coins / 2) points from node u. The number of halvings for its children's subtrees becomes h + 1. We recursively call solve(v, u, h + 1) for each child v of u and sum up the results.
- The function returns max(points1, points2).
The initial call would be solve(0, -1, 0).

A crucial observation is that after a certain number of halvings, the coin value becomes 0. Since coins[i] <= 10000, which is less than 2^14, after 14 halvings, any coin value will be 0. So, we can cap the number of halvings h at a reasonable value, say 14. If h exceeds this cap, we can treat it as if it's at the cap, as further halvings won't change the coin value (it's already 0). This slightly prunes the recursion tree but doesn't change the fundamental exponential complexity.

1import java.util.*;2 3class Solution {4    private List<Integer>[] adj;5    private int[] coins;6    private int k;7    private final int MAX_HALVINGS = 14; // log2(10000) is approx 13.3, so 14 is safe8 9    public int maximumPoints(int[][] edges, int[] coins, int k) {10        int n = coins.length;11        this.adj = new ArrayList[n];12        for (int i = 0; i < n; i++) {13            adj[i] = new ArrayList<>();14        }15        for (int[] edge : edges) {16            adj[edge[0]].add(edge[1]);17            adj[edge[1]].add(edge[0]);18        }19        this.coins = coins;20        this.k = k;21        22        return (int) solve(0, -1, 0);23    }24 25    private long solve(int u, int parent, int h) {26        // Cap the number of halvings27        if (h >= MAX_HALVINGS) {28            // After enough halvings, coins are 0. 29            // Option 1: -k points. Option 2: 0 points. Max is 0 (since k>=0).30            // This applies to the whole subtree, so total points from here is 0.31            return 0;32        }33 34        long current_coins = coins[u] >> h;35 36        // Option 1: Collect coins[u] - k37        long points1 = current_coins - k;38        for (int v : adj[u]) {39            if (v != parent) {40                points1 += solve(v, u, h);41            }42        }43 44        // Option 2: Collect floor(coins[u] / 2)45        long points2 = current_coins / 2;46        for (int v : adj[u]) {47            if (v != parent) {48                points2 += solve(v, u, h + 1);49            }50        }51 52        return Math.max(points1, points2);53    }54}

Complexity

Time

O(2^N) in the worst case. At each node, the recursion branches into two possibilities. For a tree of depth `N`, this can lead to an exponential number of calls. Even with capping the number of halvings, the number of calls is very large without memoization.

Space

O(N) for the recursion stack in the worst case (for a skewed tree) and for storing the adjacency list.

Trade-offs

Pros

Simple to understand and implement.
Correctly models the problem's state transitions.

Cons

Extremely inefficient due to recomputing solutions for the same subproblems ((node, h) pairs).
Leads to a Time Limit Exceeded (TLE) error on larger test cases due to its exponential time complexity.

The brute-force approach is inefficient because it repeatedly solves the same subproblems. We can optimize this by using memoization, a form of dynamic programming. We'll store the results of our recursive function calls in a 2D array and reuse them when the same subproblem is encountered again.

The state of our DP is defined by (u, h), representing the maximum points obtainable from the subtree rooted at node u, given that its coin value (and all coins in its subtree) has been halved h times by its ancestors.

Algorithm

Build an adjacency list for the tree.
Initialize a memoization table memo[n][MAX_HALVINGS] with a value indicating 'not computed'.
Define a recursive function solve(u, parent, h):
- If h >= MAX_HALVINGS, return 0, as coin values will be 0.
- If memo[u][h] is already computed, return it.
- Calculate current_coins = coins[u] >> h.
- Calculate points1 = (current_coins - k) + sum(solve(v, u, h)) for all children v.
- Calculate points2 = (current_coins / 2) + sum(solve(v, u, h + 1)) for all children v.
- Store max(points1, points2) in memo[u][h] and return it.
Start the recursion with solve(0, -1, 0).

Walkthrough

This approach enhances the recursive solution by adding a memoization table, memo[u][h], to store the result of solve(u, parent, h). This avoids redundant computations and drastically improves performance.

The logic is almost identical to the brute-force approach, with the addition of memoization:

Build the adjacency list for the tree.
Initialize a 2D array memo of size n x MAX_HALVINGS with a special value (e.g., null) to indicate that a state has not been computed yet. MAX_HALVINGS can be set to a small constant like 14, since log2(10000) < 14. After 14 halvings, any coin value becomes 0.
The recursive function solve(u, parent, h):
- First, check if the result for state (u, h) is already in our memoization table. If memo[u][h] is not null, return the stored result.
- Cap the number of halvings h. If h is too large, the points from the entire subtree will be 0 (since k >= 0).
- Calculate the current coin value for node u: current_coins = coins[u] >> h.
- Option 1: Calculate points for taking current_coins - k. Recursively call solve(v, u, h) for all children v.
- Option 2: Calculate points for taking floor(current_coins / 2). Recursively call solve(v, u, h + 1) for all children v.
- Store the result max(points1, points2) in memo[u][h] before returning it.
The initial call is solve(0, -1, 0).

This top-down DP approach explores the problem using a post-order traversal (DFS), calculating the optimal points for subtrees before calculating it for the parent.

1import java.util.*;2 3class Solution {4    private List<Integer>[] adj;5    private int[] coins;6    private int k;7    private Long[][] memo;8    private final int MAX_HALVINGS = 14; // log2(10000) is approx 13.3, so 14 is safe9 10    public int maximumPoints(int[][] edges, int[] coins, int k) {11        int n = coins.length;12        this.adj = new ArrayList[n];13        for (int i = 0; i < n; i++) {14            adj[i] = new ArrayList<>();15        }16        for (int[] edge : edges) {17            adj[edge[0]].add(edge[1]);18            adj[edge[1]].add(edge[0]);19        }20        this.coins = coins;21        this.k = k;22        this.memo = new Long[n][MAX_HALVINGS];23        24        return (int) solve(0, -1, 0);25    }26 27    private long solve(int u, int parent, int h) {28        if (h >= MAX_HALVINGS) {29            return 0;30        }31        32        if (memo[u][h] != null) {33            return memo[u][h];34        }35 36        long current_coins = coins[u] >> h;37 38        // Option 1: Collect coins[u] - k39        long points1 = current_coins - k;40        for (int v : adj[u]) {41            if (v != parent) {42                points1 += solve(v, u, h);43            }44        }45 46        // Option 2: Collect floor(coins[u] / 2)47        long points2 = current_coins / 2;48        for (int v : adj[u]) {49            if (v != parent) {50                points2 += solve(v, u, h + 1);51            }52        }53 54        return memo[u][h] = Math.max(points1, points2);55    }56}

Complexity

Time

O(N * log(max(coins))). The number of states is `N * MAX_HALVINGS`. Each state `(u, h)` is computed once. To compute a state, we iterate over the children of `u`. The total work across all states is proportional to the sum of degrees of all nodes, which is `O(N)`. So, the total time is `O(N * MAX_HALVINGS)`. Since `MAX_HALVINGS` is related to `log(max(coins))`, the complexity is `O(N * log(max(coins)))`.

Space

O(N * log(max(coins))). This is dominated by the memoization table `memo` of size `N * MAX_HALVINGS`. The recursion stack depth can go up to `O(N)` in the worst case (a skewed tree).

Trade-offs

Pros

Highly efficient and guaranteed to pass within time limits.
Solves the problem by breaking it down into a limited number of unique subproblems.
The use of memoization avoids exponential complexity.

Cons

Requires extra space for the memoization table.
The concept might be slightly more complex to grasp than simple recursion.

Solutions

Solution

1class Solution {2private3  int k;4private5  int[] coins;6private7  Integer[][] f;8private9  List<Integer>[] g;10public11  int maximumPoints(int[][] edges, int[] coins, int k) {12    this.k = k;13    this.coins = coins;14    int n = coins.length;15    f = new Integer[n][15];16    g = new List[n];17    Arrays.setAll(g, i->new ArrayList<>());18    for (var e : edges) {19      int a = e[0], b = e[1];20      g[a].add(b);21      g[b].add(a);22    }23    return dfs(0, -1, 0);24  }25private26  int dfs(int i, int fa, int j) {27    if (f[i][j] != null) {28      return f[i][j];29    }30    int a = (coins[i] >> j) - k;31    int b = coins[i] >> (j + 1);32    for (int c : g[i]) {33      if (c != fa) {34        a += dfs(c, i, j);35        if (j < 14) {36          b += dfs(c, i, j + 1);37        }38      }39    }40    return f[i][j] = Math.max(a, b);41  }42}

1class Solution {2public:3  int maximumPoints(vector<vector<int>> &edges, vector<int> &coins, int k) {4    int n = coins.size();5    int f[n][15];6    memset(f, -1, sizeof(f));7    vector<int> g[n];8    for (auto &e : edges) {9      int a = e[0], b = e[1];10      g[a].emplace_back(b);11      g[b].emplace_back(a);12    }13    function<int(int, int, int)> dfs = [&](int i, int fa, int j) {14      if (f[i][j] != -1) {15        return f[i][j];16      }17      int a = (coins[i] >> j) - k;18      int b = coins[i] >> (j + 1);19      for (int c : g[i]) {20        if (c != fa) {21          a += dfs(c, i, j);22          if (j < 14) {23            b += dfs(c, i, j + 1);24          }25        }26      }27      return f[i][j] = max(a, b);28    };29    return dfs(0, -1, 0);30  }31};

1class Solution:2    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

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem