Maximize Sum of Weights after Edge Removals

Hard

#2992Time: O(2^N * N). There are `n-1` edges, leading to `2^(n-1)` subsets. For each subset, we validate the degrees of `N` nodes, taking `O(N)` time.Space: O(N) for the recursion stack depth and the `degrees` array.1 company

Patterns

Algorithms

Data structures

Companies

Prompt

There exists an undirected tree with n nodes numbered 0 to n - 1. You are given a 2D integer array edges of length n - 1, where edges[i] = [u_i, v_i, w_i] indicates that there is an edge between nodes u_i and v_i with weight w_i in the tree.

Your task is to remove zero or more edges such that:

Each node has an edge with at most k other nodes, where k is given.
The sum of the weights of the remaining edges is maximized.

Return the maximum possible sum of weights for the remaining edges after making the necessary removals.

Example 1:

Input: edges = [[0,1,4],[0,2,2],[2,3,12],[2,4,6]], k = 2

Output: 22

Explanation:

Node 2 has edges with 3 other nodes. We remove the edge [0, 2, 2], ensuring that no node has edges with more than k = 2 nodes.
The sum of weights is 22, and we can't achieve a greater sum. Thus, the answer is 22.

Example 2:

Input: edges = [[0,1,5],[1,2,10],[0,3,15],[3,4,20],[3,5,5],[0,6,10]], k = 3

Output: 65

Explanation:

Since no node has edges connecting it to more than k = 3 nodes, we don't remove any edges.
The sum of weights is 65. Thus, the answer is 65.

Constraints:

2 <= n <= 10⁵
1 <= k <= n - 1
edges.length == n - 1
edges[i].length == 3
0 <= edges[i][0] <= n - 1
0 <= edges[i][1] <= n - 1
1 <= edges[i][2] <= 10⁶
The input is generated such that edges form a valid tree.

Approaches

3 approaches with complexity analysis and trade-offs.

This approach explores every possible subset of edges. For each subset, it checks if the degree constraint (each node connected to at most k other nodes) is satisfied for all nodes. If the constraint holds, the sum of weights of the edges in the subset is calculated and compared with the maximum sum found so far. This exhaustive search is implemented using recursion with backtracking.

Algorithm

Define a recursive function backtrack(edgeIndex, currentWeight, degrees).
The base case for the recursion is when edgeIndex reaches the total number of edges. In the base case, check if the degrees array for all nodes satisfies the constraint (degree <= k). If it does, update the global maximum weight. If not, this path is invalid.
In the recursive step, for the edge at edgeIndex, explore two branches:
1. Exclude the edge: Make a recursive call backtrack(edgeIndex + 1, currentWeight, degrees).
2. Include the edge: Add its weight to currentWeight, increment the degrees of its two endpoint nodes, and then make a recursive call backtrack(edgeIndex + 1, newWeight, newDegrees). Remember to backtrack by undoing the changes to degrees after the call returns.
Initialize the process by calling backtrack(0, 0, new int[n]).

Walkthrough

The brute-force method systematically generates all 2^(n-1) subsets of edges. For each subset, we form a graph and verify if the degree of every node is at most k. If this condition is met, we compute the sum of weights of the edges in the current subset and update our answer with the maximum sum found.

This can be implemented with a recursive backtracking function. The function would iterate through each edge, making a decision to either include it in our set or exclude it. After considering all edges, we have a complete subset, which we then validate against the degree constraints.

1class Solution {2    long maxWeight = 0;3    int[][] edges;4    int n;5    int k;6 7    public long maximumValueSum(int[][] edges, int k) {8        this.edges = edges;9        this.n = edges.length + 1;10        this.k = k;11        12        backtrack(0, 0, new int[n]);13        return maxWeight;14    }15 16    private void backtrack(int edgeIndex, long currentWeight, int[] degrees) {17        if (edgeIndex == edges.length) {18            for (int deg : degrees) {19                if (deg > k) {20                    return; // Invalid configuration21                }22            }23            maxWeight = Math.max(maxWeight, currentWeight);24            return;25        }26 27        int u = edges[edgeIndex][0];28        int v = edges[edgeIndex][1];29        int weight = edges[edgeIndex][2];30 31        // Case 1: Exclude the current edge32        backtrack(edgeIndex + 1, currentWeight, degrees);33 34        // Case 2: Include the current edge35        degrees[u]++;36        degrees[v]++;37        backtrack(edgeIndex + 1, currentWeight + weight, degrees);38        degrees[u]--; // backtrack39        degrees[v]--; // backtrack40    }41}

Complexity

Time

O(2^N * N). There are `n-1` edges, leading to `2^(n-1)` subsets. For each subset, we validate the degrees of `N` nodes, taking `O(N)` time.

Space

O(N) for the recursion stack depth and the `degrees` array.

Trade-offs

Pros

Simple to conceptualize and implement.

Cons

Extremely inefficient due to its exponential time complexity.
Will result in a 'Time Limit Exceeded' error for the given constraints.

This approach leverages the tree structure of the graph by using dynamic programming. A post-order traversal (DFS) is employed to compute the solution for subtrees. For each node u, we calculate the maximum possible weight sum for the subtree rooted at u under two distinct scenarios: one where the edge connecting u to its parent is kept, and another where it's dropped. This allows decisions to be propagated up the tree correctly.

Algorithm

Build an adjacency list for the tree.
Start a DFS from an arbitrary root (e.g., node 0). The function dfs(u, parent) will perform a post-order traversal.
For each node u, dfs(u, parent) returns a pair of values: [max_sum_if_parent_edge_kept, max_sum_if_parent_edge_dropped].
In the post-order step for u (after visiting all its children):
- Initialize a base_sum by summing up the dropped values from all children's results.
- For each child v, calculate the 'gain' of keeping the edge (u, v): gain = weight(u, v) + kept_v - dropped_v.
- Collect all these gains in a list.
- Sort the child_gains list in descending order.
- Calculate res_dropped: base_sum + sum of the top k positive gains.
- Calculate res_kept: base_sum + sum of the top k-1 positive gains.
The final answer is the dropped value returned by the initial dfs(root, -1) call.

Walkthrough

We first represent the tree using an adjacency list. Then, we perform a DFS from an arbitrary root (e.g., node 0). The core of this method is a function, dfs(u, parent), which computes results for the subtree at u after all its children's subtrees have been processed.

This function returns two values for node u:

res_kept: The maximum weight sum in u's subtree, assuming the edge to its parent p is kept.
res_dropped: The maximum weight sum in u's subtree, assuming the edge to p is dropped.

To compute these, we first recursively call dfs on all children v of u. We aggregate the dropped results from all children into a base_sum. Then, for each child v, we calculate the potential 'gain' from keeping the edge (u, v), which is weight(u, v) + v_kept - v_dropped. We collect these gains, sort them in descending order, and add the top positive gains to base_sum based on the degree budget at u. If the parent edge is dropped, u can connect to k children. If it's kept, it can connect to k-1 children.

1class Solution {2    private List<List<int[]>> adj;3    private int k;4 5    public long maximumValueSum(int[][] edges, int k) {6        int n = edges.length + 1;7        this.adj = new ArrayList<>(n);8        for (int i = 0; i < n; i++) {9            adj.add(new ArrayList<>());10        }11        for (int[] edge : edges) {12            adj.get(edge[0]).add(new int[]{edge[1], edge[2]});13            adj.get(edge[1]).add(new int[]{edge[0], edge[2]});14        }15        this.k = k;16        17        long[] result = dfs(0, -1);18        return result[1]; // [0] is kept, [1] is dropped19    }20 21    // Returns a pair: [max_sum_if_parent_edge_kept, max_sum_if_parent_edge_dropped]22    private long[] dfs(int u, int p) {23        List<Long> childGains = new ArrayList<>();24        long baseSumFromChildren = 0;25 26        for (int[] edgeInfo : adj.get(u)) {27            int v = edgeInfo[0];28            int weight = edgeInfo[1];29            if (v == p) continue;30 31            long[] childResult = dfs(v, u);32            long keptV = childResult[0];33            long droppedV = childResult[1];34            35            baseSumFromChildren += droppedV;36            childGains.add((long)weight + keptV - droppedV);37        }38 39        Collections.sort(childGains, Collections.reverseOrder());40 41        long resKept = baseSumFromChildren;42        long resDropped = baseSumFromChildren;43 44        for (int i = 0; i < childGains.size(); i++) {45            long gain = childGains.get(i);46            if (gain <= 0) break;47            48            if (i < k) { // Can keep up to k edges if parent edge is dropped49                resDropped += gain;50            }51            if (i < k - 1) { // Can keep up to k-1 edges if parent edge is kept52                resKept += gain;53            }54        }55        56        return new long[]{resKept, resDropped};57    }58}

Complexity

Time

O(N log N). The DFS visits each node and edge once. At a node `u` with degree `d`, sorting the `d-1` child gains takes `O(d log d)`. The sum of `d log d` over all nodes is bounded by `O(N log N)`.

Space

O(N). Required for the adjacency list, the recursion stack (which can go up to `O(N)` deep), and the temporary list for child gains (which can be up to `O(N)` for a star graph).

Trade-offs

Pros

Efficient and correct for the given constraints.
Systematically solves the problem by breaking it down into subproblems on the tree.

Cons

The sorting step at each node can be a bottleneck for nodes with a very high degree, making it slightly less optimal than the heap-based approach.

This approach is an optimization of the Tree DP method. Instead of sorting the list of child gains at each node, which takes O(d log d) time for a node with d children, we use a min-heap to find the top k gains more efficiently. This reduces the time spent at each node from O(d log d) to O(d log k), leading to a better overall time complexity.

Algorithm

The overall DP state and traversal logic are the same as the sorting-based approach.
The key change is in how the top gains are selected at each node u.
Instead of collecting all child gains and sorting, we process them as they are computed.
We maintain a min-heap of size at most k and a running sum of the gains in the heap.
For each child v, calculate its gain. If the gain is positive:
- Add the gain to the heap and the running sum.
- If the heap size exceeds k, remove the smallest element (the root) from the heap and subtract it from the running sum.
After processing all children, the running sum will be the total of the top k positive gains (gain_sum_for_dropped).
To find the sum of the top k-1 gains (gain_sum_for_kept), we check if the heap size is exactly k. If so, we subtract the smallest element (the heap's root) from gain_sum_for_dropped.

Walkthrough

The fundamental dynamic programming structure remains the same. We still perform a post-order traversal and compute (kept, dropped) values for each subtree. The optimization lies in avoiding the O(d log d) sorting step at each node.

At a node u, as we compute the gain from each child v, we use a min-heap of size k to maintain the k largest gains encountered so far. For each positive gain, we add it to the heap. If the heap's size grows beyond k, we evict the smallest element. This process takes O(log k) time per child.

After iterating through all children of u, the heap contains the top k positive gains. The sum of these gains is the value we add to base_sum for the res_dropped case. For the res_kept case, we need the sum of the top k-1 gains. If the heap's size is k, this is simply the total sum minus the smallest element in the heap (the root). If the heap's size is less than k, the sum for both cases is the same.

1class Solution {2    private List<List<int[]>> adj;3    private int k;4 5    public long maximumValueSum(int[][] edges, int k) {6        int n = edges.length + 1;7        this.adj = new ArrayList<>(n);8        for (int i = 0; i < n; i++) {9            adj.add(new ArrayList<>());10        }11        for (int[] edge : edges) {12            adj.get(edge[0]).add(new int[]{edge[1], edge[2]});13            adj.get(edge[1]).add(new int[]{edge[0], edge[2]});14        }15        this.k = k;16        17        long[] result = dfs(0, -1);18        return result[1];19    }20 21    private long[] dfs(int u, int p) {22        long baseSumFromChildren = 0;23        24        PriorityQueue<Long> topGainsHeap = new PriorityQueue<>();25        long topGainsSum = 0;26 27        for (int[] edgeInfo : adj.get(u)) {28            int v = edgeInfo[0];29            if (v == p) continue;30            31            int weight = edgeInfo[1];32            long[] childResult = dfs(v, u);33            long keptV = childResult[0];34            long droppedV = childResult[1];35            36            baseSumFromChildren += droppedV;37            long gain = (long)weight + keptV - droppedV;38 39            if (gain > 0) {40                topGainsHeap.add(gain);41                topGainsSum += gain;42                if (topGainsHeap.size() > k) {43                    topGainsSum -= topGainsHeap.poll();44                }45            }46        }47 48        long resDropped = baseSumFromChildren + topGainsSum;49        long resKept = resDropped;50        51        if (topGainsHeap.size() == k) {52            resKept -= topGainsHeap.peek();53        }54        55        return new long[]{resKept, resDropped};56    }57}

Complexity

Time

O(N log k). The DFS visits each node once. At each node `u` with `d` children, we perform `d` heap operations, each taking `O(log k)` time. The total time is `Sum(d * log k)` over all nodes, which simplifies to `O(N log k)`.

Space

O(N). For the adjacency list, recursion stack, and the heaps. The total size of all heaps on the call stack is bounded by `O(N)`.

Trade-offs

Pros

The most efficient solution for this problem.
Improves upon the sorting approach, especially when k is small compared to N.

Cons

The implementation can be slightly more complex than the sorting approach.

Solutions

Solution

1class Solution {2private3  List<int[]>[] g;4private5  int k;6public7  long maximizeSumOfWeights(int[][] edges, int k) {8    this.k = k;9    int n = edges.length + 1;10    g = new List[n];11    Arrays.setAll(g, i->new ArrayList<>());12    for (var e : edges) {13      int u = e[0], v = e[1], w = e[2];14      g[u].add(new int[]{v, w});15      g[v].add(new int[]{u, w});16    }17    var ans = dfs(0, -1);18    return Math.max(ans[0], ans[1]);19  }20private21  long[] dfs(int u, int fa) {22    long s = 0;23    List<Long> t = new ArrayList<>();24    for (var e : g[u]) {25      int v = e[0], w = e[1];26      if (v == fa) {27        continue;28      }29      var res = dfs(v, u);30      s += res[0];31      long d = w + res[1] - res[0];32      if (d > 0) {33        t.add(d);34      }35    }36    t.sort(Comparator.reverseOrder());37    for (int i = 0; i < Math.min(t.size(), k - 1); ++i) {38      s += t.get(i);39    }40    return new long[]{s + (t.size() >= k ? t.get(k - 1) : 0), s};41  }42}

1class Solution {2public:3  long long maximizeSumOfWeights(vector<vector<int>> &edges, int k) {4    int n = edges.size() + 1;5    vector<vector<pair<int, int>>> g(n);6    for (auto &e : edges) {7      int u = e[0], v = e[1], w = e[2];8      g[u].emplace_back(v, w);9      g[v].emplace_back(u, w);10    }11    using ll = long long;12    auto dfs = [&](auto &&dfs, int u, int fa) -> pair<ll, ll> {13      ll s = 0;14      vector<ll> t;15      for (auto &[v, w] : g[u]) {16        if (v == fa) {17          continue;18        }19        auto [a, b] = dfs(dfs, v, u);20        s += a;21        ll d = w + b - a;22        if (d > 0) {23          t.push_back(d);24        }25      }26      ranges ::sort(t, greater<>());27      for (int i = 0; i < min((int)t.size(), k - 1); ++i) {28        s += t[i];29      }30      return {s + (t.size() >= k ? t[k - 1] : 0), s};31    };32    auto [x, y] = dfs(dfs, 0, -1);33    return max(x, y);34  }35};

1class Solution:2    def maximizeSumOfWeights(self, edges: List[List[int]], k: int) -> int: def dfs(u: int, fa: int) -> Tuple[int, int]: s = 0 t = [] for v, w in g[u]: if v == fa: continue a, b = dfs(v, u) s += a if (d: = (w + b - a)) > 0: t . append(d) t . sort(reverse=True) return s + sum(t[: k]), s + sum(t[: k - 1]) n = len(edges) + 1 g: List[List[Tuple[int, int]]] = [[] for _ in range(n)] for u, v, w in edges: g[u]. append((v, w)) g[v]. append((u, w)) x, y = dfs(0, - 1) return max(x, y)

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

Solution

1class Solution {2private3  List<int[]>[] g;4private5  int k;6public7  long maximizeSumOfWeights(int[][] edges, int k) {8    this.k = k;9    int n = edges.length + 1;10    g = new List[n];11    Arrays.setAll(g, i->new ArrayList<>());12    for (var e : edges) {13      int u = e[0], v = e[1], w = e[2];14      g[u].add(new int[]{v, w});15      g[v].add(new int[]{u, w});16    }17    var ans = dfs(0, -1);18    return Math.max(ans[0], ans[1]);19  }20private21  long[] dfs(int u, int fa) {22    long s = 0;23    List<Long> t = new ArrayList<>();24    for (var e : g[u]) {25      int v = e[0], w = e[1];26      if (v == fa) {27        continue;28      }29      var res = dfs(v, u);30      s += res[0];31      long d = w + res[1] - res[0];32      if (d > 0) {33        t.add(d);34      }35    }36    t.sort(Comparator.reverseOrder());37    for (int i = 0; i < Math.min(t.size(), k - 1); ++i) {38      s += t.get(i);39    }40    return new long[]{s + (t.size() >= k ? t.get(k - 1) : 0), s};41  }42}

1class Solution {2public:3  long long maximizeSumOfWeights(vector<vector<int>> &edges, int k) {4    int n = edges.size() + 1;5    vector<vector<pair<int, int>>> g(n);6    for (auto &e : edges) {7      int u = e[0], v = e[1], w = e[2];8      g[u].emplace_back(v, w);9      g[v].emplace_back(u, w);10    }11    using ll = long long;12    auto dfs = [&](auto &&dfs, int u, int fa) -> pair<ll, ll> {13      ll s = 0;14      vector<ll> t;15      for (auto &[v, w] : g[u]) {16        if (v == fa) {17          continue;18        }19        auto [a, b] = dfs(dfs, v, u);20        s += a;21        ll d = w + b - a;22        if (d > 0) {23          t.push_back(d);24        }25      }26      ranges ::sort(t, greater<>());27      for (int i = 0; i < min((int)t.size(), k - 1); ++i) {28        s += t[i];29      }30      return {s + (t.size() >= k ? t[k - 1] : 0), s};31    };32    auto [x, y] = dfs(dfs, 0, -1);33    return max(x, y);34  }35};

1class Solution:2    def maximizeSumOfWeights(self, edges: List[List[int]], k: int) -> int: def dfs(u: int, fa: int) -> Tuple[int, int]: s = 0 t = [] for v, w in g[u]: if v == fa: continue a, b = dfs(v, u) s += a if (d: = (w + b - a)) > 0: t . append(d) t . sort(reverse=True) return s + sum(t[: k]), s + sum(t[: k - 1]) n = len(edges) + 1 g: List[List[Tuple[int, int]]] = [[] for _ in range(n)] for u, v, w in edges: g[u]. append((v, w)) g[v]. append((u, w)) x, y = dfs(0, - 1) return max(x, y)

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.