LeetCode 3367 - Maximize Sum of Weights after Edge Removals

Difficulty: 🔴 Hard
Topics: Dynamic Programming, Tree, Depth-First Search, Sorting

Solution

Problem Understanding

We are given an undirected weighted tree with n nodes and n - 1 edges. Since the graph is a tree, there is exactly one simple path between any pair of nodes, and there are no cycles.

Each edge has a weight, and we may remove any number of edges. After removals, every node must have degree at most k. In other words, no node is allowed to remain connected to more than k neighboring nodes.

The goal is to maximize the total weight of all remaining edges.

The input is represented as:

edges[i] = [u, v, w]
u and v are endpoints of an edge
w is the edge weight
k is the maximum allowed degree for every node

The output is a single integer, the largest possible sum of weights after removing edges while respecting the degree constraint.

The constraints are extremely important:

n can be as large as 10^5
Edge weights can be as large as 10^6

A brute-force solution that tries all subsets of edges is completely impossible because there are 2^(n-1) possible subsets.

The fact that the graph is a tree is the key observation. Trees have no cycles, which allows dynamic programming with depth-first search.

Several edge cases are important:

k >= maximum degree of the tree, in which case no edge needs to be removed.
k = 1, where the result becomes a maximum-weight matching on a tree.
A star-shaped tree where one node connects to many neighbors, forcing many removals.
Very deep trees, where recursive DFS implementations may hit recursion limits in Python.
Situations where keeping a lower-weight edge enables better overall choices deeper in the subtree.

The problem guarantees that the input graph is always a valid tree, so connectivity and cycle validation are unnecessary.

Approaches

Brute Force Approach

The most direct idea is to consider every subset of edges:

For each subset, compute the degree of every node.
Check whether all degrees are at most k.
If valid, compute the sum of remaining edge weights.
Track the maximum sum.

This approach is guaranteed to find the optimal answer because it explores every possible valid configuration.

However, a tree with n nodes has n - 1 edges, so there are:

$$2^{n-1}$$

possible subsets.

With n = 10^5, this is astronomically large and completely infeasible.

Optimal Dynamic Programming on Trees

The important observation is that the graph is a tree. When rooted at any node, every subtree becomes independent except for the edge connecting it to its parent.

This naturally suggests tree DP.

For each node, we want to know:

What is the best result if this node still has one degree slot available for its parent edge?
What is the best result if all k degree slots can be used by child edges?

This creates two DP states per node.

For every child edge, we compute the profit gained by keeping that edge instead of removing it. Some edges contribute positively, others negatively.

Then the problem becomes:

Select up to some number of child edges with the largest positive gains.

This transforms the difficult global optimization problem into a local greedy selection after DFS computation.

Approach Comparison

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(2^n * n)	O(n)	Tries every subset of edges
Optimal	O(n log n)	O(n)	Tree DP with greedy selection of best child contributions

Algorithm Walkthrough

Step 1: Build the adjacency list

Because the graph is a tree, an adjacency list is the most efficient representation.

For every edge [u, v, w]:

Add (v, w) to u
Add (u, w) to v

This gives efficient traversal of neighbors.

Step 2: Define the DP states

For each node u, we compute two values:

dp0[u]
Maximum weight achievable in the subtree rooted at u
Assuming u may still use all k connections toward children
dp1[u]
Maximum weight achievable in the subtree rooted at u
Assuming one degree slot is already consumed by the parent edge
Therefore only k - 1 child edges may remain

The distinction is crucial because parent-child relationships consume degree capacity.

Step 3: DFS traversal

Run DFS from an arbitrary root, usually node 0.

For every child v of node u:

Recursively compute DP values for v
Initially assume the edge (u, v) is removed
The baseline contribution is dp0[v]

Now consider keeping edge (u, v).

If we keep it:

We gain edge weight w
The child v must use one degree slot for its parent
Therefore we use dp1[v]

The improvement from keeping the edge is:

$$gain = w + dp1[v] - dp0[v]$$

Step 4: Collect profitable gains

Only positive gains are useful.

If gain <= 0, removing the edge is always better.

Store all positive gains in an array.

Step 5: Greedily select best gains

Sort gains in descending order.

For dp0[u]:

We may keep at most k child edges

So we add the largest k gains.

For dp1[u]:

One slot is reserved for the parent edge
We may keep at most k - 1 child edges

So we add the largest k - 1 gains.

Step 6: Return the root result

The root has no parent, so it uses the unrestricted state:

$$dp0[root]$$

This is the final answer.

Why it works

The DP works because subtrees in a tree are independent once the parent-child relationship is fixed.

For each child edge, there are only two possibilities:

Remove it
Keep it

The gain formula precisely measures whether keeping the edge improves the optimal subtree result.

Since child decisions become independent after DFS computation, selecting the best positive gains greedily is optimal. Choosing the largest gains maximizes the total contribution while respecting the degree limit.

Python Solution

from typing import List
import sys

sys.setrecursionlimit(1 << 20)

class Solution:
    def maximizeSumOfWeights(self, edges: List[List[int]], k: int) -> int:
        n = len(edges) + 1

        graph = [[] for _ in range(n)]

        for u, v, w in edges:
            graph[u].append((v, w))
            graph[v].append((u, w))

        def dfs(node: int, parent: int):
            base = 0
            gains = []

            for neighbor, weight in graph[node]:
                if neighbor == parent:
                    continue

                child_dp0, child_dp1 = dfs(neighbor, node)

                base += child_dp0

                gain = weight + child_dp1 - child_dp0

                if gain > 0:
                    gains.append(gain)

            gains.sort(reverse=True)

            dp0 = base
            dp1 = base

            for i in range(min(k, len(gains))):
                dp0 += gains[i]

            for i in range(min(k - 1, len(gains))):
                dp1 += gains[i]

            return dp0, dp1

        answer, _ = dfs(0, -1)
        return answer

The implementation begins by constructing an adjacency list for the tree.

The DFS function returns two DP values:

dp0, where the current node may still use all k degree slots for child edges
dp1, where one slot is already consumed by the parent edge

For each child:

We first assume the edge is removed
Then compute the benefit of keeping it

The gain formula:

gain = weight + child_dp1 - child_dp0

represents the net improvement from preserving the edge.

Only positive gains are useful, so negative values are ignored.

After sorting gains in descending order:

dp0 takes the best k
dp1 takes the best k - 1

Finally, the DFS result for the root is returned.

Go Solution

package main

import (
	"sort"
)

func maximizeSumOfWeights(edges [][]int, k int) int64 {
	n := len(edges) + 1

	type Edge struct {
		to     int
		weight int
	}

	graph := make([][]Edge, n)

	for _, e := range edges {
		u, v, w := e[0], e[1], e[2]

		graph[u] = append(graph[u], Edge{v, w})
		graph[v] = append(graph[v], Edge{u, w})
	}

	var dfs func(int, int) (int64, int64)

	dfs = func(node int, parent int) (int64, int64) {
		var base int64 = 0
		gains := []int64{}

		for _, edge := range graph[node] {
			neighbor := edge.to
			weight := edge.weight

			if neighbor == parent {
				continue
			}

			childDP0, childDP1 := dfs(neighbor, node)

			base += childDP0

			gain := int64(weight) + childDP1 - childDP0

			if gain > 0 {
				gains = append(gains, gain)
			}
		}

		sort.Slice(gains, func(i, j int) bool {
			return gains[i] > gains[j]
		})

		dp0 := base
		dp1 := base

		limit0 := min(k, len(gains))
		limit1 := min(k-1, len(gains))

		for i := 0; i < limit0; i++ {
			dp0 += gains[i]
		}

		for i := 0; i < limit1; i++ {
			dp1 += gains[i]
		}

		return dp0, dp1
	}

	answer, _ := dfs(0, -1)

	return answer
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

The Go implementation follows exactly the same DP logic as the Python version.

Several Go-specific details are important:

int64 is used because the total edge weight can exceed 32-bit integer range.
Slices are used for adjacency lists and gain storage.
sort.Slice sorts gains in descending order.
Recursive closures require explicit declaration using var dfs func(...).

Worked Examples

Example 1

Input:

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

Tree Structure

DFS Processing

Leaf Nodes

Nodes 1, 3, and 4 are leaves.

Node	dp0	dp1
1	0	0
3	0	0
4	0	0

Process Node 2

Children:

(2,3) weight 12
(2,4) weight 6

Base:

0 + 0 = 0

Gains:

Edge	Gain
(2,3)	12 + 0 - 0 = 12
(2,4)	6 + 0 - 0 = 6

Sorted gains:

[12, 6]

Since k = 2:

dp0[2] = 0 + 12 + 6 = 18
dp1[2] = 0 + 12 = 12

Process Node 0

Children:

(0,1) weight 4
(0,2) weight 2

Base:

0 + 18 = 18

Gains:

Edge	Gain
(0,1)	4 + 0 - 0 = 4
(0,2)	2 + 12 - 18 = -4

Only positive gains:

[4]

Final:

dp0[0] = 18 + 4 = 22

Answer:

Example 2

Input:

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

No node exceeds degree 3.

Every gain remains positive, so all edges are preserved.

Total:

5 + 10 + 15 + 20 + 5 + 10 = 65

Answer:

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n log n)	Each node sorts its gain list
Space	O(n)	Adjacency list, recursion stack, and gain storage

The DFS visits every edge exactly once, contributing O(n) work.

The dominant cost comes from sorting gain arrays. Across the entire tree, the total number of gains equals n - 1. In the worst case, one node may have degree O(n), giving overall complexity O(n log n).

The space complexity is linear because the adjacency list stores all edges once in each direction, and recursion depth may reach O(n) in a skewed tree.

Test Cases

from typing import List

class Solution:
    def maximizeSumOfWeights(self, edges: List[List[int]], k: int) -> int:
        import sys
        sys.setrecursionlimit(1 << 20)

        n = len(edges) + 1

        graph = [[] for _ in range(n)]

        for u, v, w in edges:
            graph[u].append((v, w))
            graph[v].append((u, w))

        def dfs(node, parent):
            base = 0
            gains = []

            for neighbor, weight in graph[node]:
                if neighbor == parent:
                    continue

                child0, child1 = dfs(neighbor, node)

                base += child0

                gain = weight + child1 - child0

                if gain > 0:
                    gains.append(gain)

            gains.sort(reverse=True)

            dp0 = base
            dp1 = base

            for i in range(min(k, len(gains))):
                dp0 += gains[i]

            for i in range(min(k - 1, len(gains))):
                dp1 += gains[i]

            return dp0, dp1

        return dfs(0, -1)[0]

sol = Solution()

# Example 1
assert sol.maximizeSumOfWeights(
    [[0,1,4],[0,2,2],[2,3,12],[2,4,6]], 2
) == 22  # must remove one edge from node 2

# Example 2
assert sol.maximizeSumOfWeights(
    [[0,1,5],[1,2,10],[0,3,15],[3,4,20],[3,5,5],[0,6,10]], 3
) == 65  # no removals needed

# Smallest tree
assert sol.maximizeSumOfWeights(
    [[0,1,7]], 1
) == 7  # single edge kept

# k = 1, maximum matching behavior
assert sol.maximizeSumOfWeights(
    [[0,1,5],[1,2,6]], 1
) == 6  # choose larger edge

# Star graph
assert sol.maximizeSumOfWeights(
    [[0,1,1],[0,2,2],[0,3,3],[0,4,4]], 2
) == 7  # keep top two edges

# Chain graph
assert sol.maximizeSumOfWeights(
    [[0,1,1],[1,2,2],[2,3,3],[3,4,4]], 2
) == 10  # all edges valid

# Heavy subtree tradeoff
assert sol.maximizeSumOfWeights(
    [[0,1,1],[1,2,100],[1,3,100]], 2
) == 200  # remove parent edge

# All equal weights
assert sol.maximizeSumOfWeights(
    [[0,1,5],[0,2,5],[0,3,5]], 2
) == 10  # any two edges

# Large edge weights
assert sol.maximizeSumOfWeights(
    [[0,1,10**6],[0,2,10**6],[0,3,10**6]], 2
) == 2 * 10**6  # verifies large sums

Test Summary

Test	Why
Example 1	Basic removal scenario
Example 2	No removals needed
Smallest tree	Minimum valid input
k = 1	Matching-style constraint
Star graph	High-degree node handling
Chain graph	Deep tree with no removals
Heavy subtree tradeoff	Validates DP decisions
Equal weights	Tie-handling correctness
Large weights	Integer overflow safety

Edge Cases

Case 1: `k` Larger Than Every Node Degree

If every node already satisfies the degree constraint, the optimal solution is to keep all edges.

This case is important because an incorrect implementation might unnecessarily remove edges or mishandle gain calculations.

The DP naturally handles this because all positive gains are selected, and there are enough available degree slots to preserve every edge.

Case 2: `k = 1`

When every node can connect to at most one other node, the problem effectively becomes finding a maximum-weight matching on the tree.

This is tricky because local greedy choices can fail badly. Keeping one edge may block multiple heavier combinations elsewhere.

The two-state DP correctly handles this because each node tracks whether its parent edge already consumes its single available slot.

Case 3: Star-Shaped Tree

Consider a central node connected to many leaves.

      1
      |
2 --  0 -- 3
      |
      4

If k = 2, most edges must be removed.

A naive solution might remove arbitrary edges, but the optimal strategy is to keep the highest-weight edges.

The algorithm handles this by computing gains and selecting the top k contributions after sorting.

Case 4: Deep Linear Tree

A chain-shaped tree can produce recursion depth close to 10^5.

Without increasing Python's recursion limit, DFS may crash with a recursion depth error.

The implementation avoids this issue with:

sys.setrecursionlimit(1 << 20)

which safely supports deep recursion.