# Maximize Sum of Weights after Edge Removals
**Difficulty:** HARD
[External](https://leetcode.com/problems/maximize-sum-of-weights-after-edge-removals)
Canonical: https://scaleengineer.com/dsa/problems/maximize-sum-of-weights-after-edge-removals
**Patterns:** [Dynamic Programming](https://scaleengineer.com/dsa/patterns/dynamic-programming)
**Algorithms:** [Depth-First Search](https://scaleengineer.com/algorithms/depth-first-search)
**Data structures:** Tree
**Companies:** [Gameskraft](https://scaleengineer.com/companies/gameskraft)
---
## Problem
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] = [ui, vi, wi]` indicates that there is an edge between nodes `ui` and `vi` with weight `wi` 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:**

![](https://assets.glich.co/dsa/maximize-sum-of-weights-after-edge-removals/image0.png)

* 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 <= 105`
* `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] <= 106`
* The input is generated such that `edges` form a valid tree.

# Approaches
## Brute Force with Backtracking
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.
**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.
**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.
### Explanation
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.

```java
class Solution {
    long maxWeight = 0;
    int[][] edges;
    int n;
    int k;

    public long maximumValueSum(int[][] edges, int k) {
        this.edges = edges;
        this.n = edges.length + 1;
        this.k = k;
        
        backtrack(0, 0, new int[n]);
        return maxWeight;
    }

    private void backtrack(int edgeIndex, long currentWeight, int[] degrees) {
        if (edgeIndex == edges.length) {
            for (int deg : degrees) {
                if (deg > k) {
                    return; // Invalid configuration
                }
            }
            maxWeight = Math.max(maxWeight, currentWeight);
            return;
        }

        int u = edges[edgeIndex][0];
        int v = edges[edgeIndex][1];
        int weight = edges[edgeIndex][2];

        // Case 1: Exclude the current edge
        backtrack(edgeIndex + 1, currentWeight, degrees);

        // Case 2: Include the current edge
        degrees[u]++;
        degrees[v]++;
        backtrack(edgeIndex + 1, currentWeight + weight, degrees);
        degrees[u]--; // backtrack
        degrees[v]--; // backtrack
    }
}
```
### 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])`.

## Tree DP with Post-order Traversal and Sorting
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.
**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).
**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.
### Explanation
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`:
1.  `res_kept`: The maximum weight sum in `u`'s subtree, assuming the edge to its parent `p` is kept.
2.  `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.

```java
class Solution {
    private List<List<int[]>> adj;
    private int k;

    public long maximumValueSum(int[][] edges, int k) {
        int n = edges.length + 1;
        this.adj = new ArrayList<>(n);
        for (int i = 0; i < n; i++) {
            adj.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            adj.get(edge[0]).add(new int[]{edge[1], edge[2]});
            adj.get(edge[1]).add(new int[]{edge[0], edge[2]});
        }
        this.k = k;
        
        long[] result = dfs(0, -1);
        return result[1]; // [0] is kept, [1] is dropped
    }

    // Returns a pair: [max_sum_if_parent_edge_kept, max_sum_if_parent_edge_dropped]
    private long[] dfs(int u, int p) {
        List<Long> childGains = new ArrayList<>();
        long baseSumFromChildren = 0;

        for (int[] edgeInfo : adj.get(u)) {
            int v = edgeInfo[0];
            int weight = edgeInfo[1];
            if (v == p) continue;

            long[] childResult = dfs(v, u);
            long keptV = childResult[0];
            long droppedV = childResult[1];
            
            baseSumFromChildren += droppedV;
            childGains.add((long)weight + keptV - droppedV);
        }

        Collections.sort(childGains, Collections.reverseOrder());

        long resKept = baseSumFromChildren;
        long resDropped = baseSumFromChildren;

        for (int i = 0; i < childGains.size(); i++) {
            long gain = childGains.get(i);
            if (gain <= 0) break;
            
            if (i < k) { // Can keep up to k edges if parent edge is dropped
                resDropped += gain;
            }
            if (i < k - 1) { // Can keep up to k-1 edges if parent edge is kept
                resKept += gain;
            }
        }
        
        return new long[]{resKept, resDropped};
    }
}
```
### 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.

## Optimized Tree DP with Min-Heap
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.
**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)`.
**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.
### Explanation
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.

```java
class Solution {
    private List<List<int[]>> adj;
    private int k;

    public long maximumValueSum(int[][] edges, int k) {
        int n = edges.length + 1;
        this.adj = new ArrayList<>(n);
        for (int i = 0; i < n; i++) {
            adj.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            adj.get(edge[0]).add(new int[]{edge[1], edge[2]});
            adj.get(edge[1]).add(new int[]{edge[0], edge[2]});
        }
        this.k = k;
        
        long[] result = dfs(0, -1);
        return result[1];
    }

    private long[] dfs(int u, int p) {
        long baseSumFromChildren = 0;
        
        PriorityQueue<Long> topGainsHeap = new PriorityQueue<>();
        long topGainsSum = 0;

        for (int[] edgeInfo : adj.get(u)) {
            int v = edgeInfo[0];
            if (v == p) continue;
            
            int weight = edgeInfo[1];
            long[] childResult = dfs(v, u);
            long keptV = childResult[0];
            long droppedV = childResult[1];
            
            baseSumFromChildren += droppedV;
            long gain = (long)weight + keptV - droppedV;

            if (gain > 0) {
                topGainsHeap.add(gain);
                topGainsSum += gain;
                if (topGainsHeap.size() > k) {
                    topGainsSum -= topGainsHeap.poll();
                }
            }
        }

        long resDropped = baseSumFromChildren + topGainsSum;
        long resKept = resDropped;
        
        if (topGainsHeap.size() == k) {
            resKept -= topGainsHeap.peek();
        }
        
        return new long[]{resKept, resDropped};
    }
}
```
### 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`.

# Solutions
### Java

```java
class Solution {
private
  List<int[]>[] g;
private
  int k;
public
  long maximizeSumOfWeights(int[][] edges, int k) {
    this.k = k;
    int n = edges.length + 1;
    g = new List[n];
    Arrays.setAll(g, i->new ArrayList<>());
    for (var e : edges) {
      int u = e[0], v = e[1], w = e[2];
      g[u].add(new int[]{v, w});
      g[v].add(new int[]{u, w});
    }
    var ans = dfs(0, -1);
    return Math.max(ans[0], ans[1]);
  }
private
  long[] dfs(int u, int fa) {
    long s = 0;
    List<Long> t = new ArrayList<>();
    for (var e : g[u]) {
      int v = e[0], w = e[1];
      if (v == fa) {
        continue;
      }
      var res = dfs(v, u);
      s += res[0];
      long d = w + res[1] - res[0];
      if (d > 0) {
        t.add(d);
      }
    }
    t.sort(Comparator.reverseOrder());
    for (int i = 0; i < Math.min(t.size(), k - 1); ++i) {
      s += t.get(i);
    }
    return new long[]{s + (t.size() >= k ? t.get(k - 1) : 0), s};
  }
}

```

### CPP

```cpp
class Solution {
public:
  long long maximizeSumOfWeights(vector<vector<int>> &edges, int k) {
    int n = edges.size() + 1;
    vector<vector<pair<int, int>>> g(n);
    for (auto &e : edges) {
      int u = e[0], v = e[1], w = e[2];
      g[u].emplace_back(v, w);
      g[v].emplace_back(u, w);
    }
    using ll = long long;
    auto dfs = [&](auto &&dfs, int u, int fa) -> pair<ll, ll> {
      ll s = 0;
      vector<ll> t;
      for (auto &[v, w] : g[u]) {
        if (v == fa) {
          continue;
        }
        auto [a, b] = dfs(dfs, v, u);
        s += a;
        ll d = w + b - a;
        if (d > 0) {
          t.push_back(d);
        }
      }
      ranges ::sort(t, greater<>());
      for (int i = 0; i < min((int)t.size(), k - 1); ++i) {
        s += t[i];
      }
      return {s + (t.size() >= k ? t[k - 1] : 0), s};
    };
    auto [x, y] = dfs(dfs, 0, -1);
    return max(x, y);
  }
};

```

### Python

```python
class Solution:
    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)

```