Skip to main content
Talentd

Maximize Sum of Weights after Edge Removals - Solution & Explanation

HardDynamic ProgrammingTreeDepth-First SearchSorting11 min readAsked at: Gameskraft
Practice this problem

Problem Statement

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:

  • 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.

Approach Overview

Problem Overview: You are given a weighted tree and must remove certain edges so the remaining structure satisfies a constraint while maximizing the total weight of the kept edges. The challenge is deciding which edges contribute positive value to the final structure without violating the rule on how many connections can remain.

Approach 1: Greedy Method Using Priority Queue (O(n log n) time, O(n) space)

This approach treats the tree as a dynamic programming problem over a depth-first search. Start from any root and compute the contribution each child edge provides to the total score. For every node, calculate the "gain" from keeping a child edge versus removing it. These gains are pushed into a priority queue so you can always pick the most valuable edges first. If the node can only keep a limited number of edges, you select the top gains and discard the rest. Sorting or a heap ensures the best contributions are chosen greedily. The DFS aggregates results bottom‑up, giving an optimal selection of edges. Time complexity is O(n log n) due to heap operations, and space complexity is O(n) for recursion and storage.

Approach 2: Kruskal's Algorithm with Modifications (O(E log E) time, O(n) space)

This approach reframes the problem as building a constrained maximum spanning structure. Sort all edges by weight in descending order, similar to Kruskal’s algorithm from graph theory. While processing edges, maintain connectivity and constraint conditions using a disjoint set union structure. Edges that violate the allowed structure or exceed node constraints are skipped. Because edges are processed from highest to lowest weight, the algorithm greedily preserves the most valuable ones first. The sorting step dominates the complexity at O(E log E), while union–find operations remain near constant time with path compression. Space complexity is O(n) for the DSU arrays.

Recommended for interviews: The DFS greedy approach is typically expected. It shows you understand dynamic programming on trees and how to combine it with greedy selection using sorting or heaps. Interviewers often look for the bottom‑up DFS that computes gains and selects the best contributions. The Kruskal-style solution is useful conceptually but appears less frequently in interviews because the tree DP formulation is more direct.

Approach 1: Approach 1: Greedy Method Using Priority Queue

This approach involves using a priority queue to manage the edges we can potentially remove. By sorting the edges initially based on their weights (descending), we can prioritize keeping higher-weighted edges in the final solution. The key here is to check the degree of each node after considering an edge and maintain the constraint that the degree doesn’t exceed k.

The Python solution uses a disjoint-set (union-find) to manage which nodes are connected while tracking the degree of each node. We iterate over the sorted edges and select the ones that do not violate the constraint k. The function returns the maximum sum of edge weights while adhering to the node degree constraint.

Code

PythonC++

Python

C++

Complexity

Time Complexity: O(n log n) due to sorting and the find-union operations.

Space Complexity: O(n) for the disjoint-set data structure and degree tracking.

Try this approach in the editor →

Approach 2: Approach 2: Kruskal's Algorithm with Modifications

This approach modifies the classic Kruskal's algorithm. Instead of focusing solely on forming a minimum spanning tree, the algorithm attempts to form a tree while respecting the degree constraint by picking edges starting from the one with the highest weight.

The Java implementation adapts Kruskal's strategy to build a subset of maximal weighted edges without exceeding the node degree limitations. By processing edges from the highest weight to the lowest, it effectively manages to keep the subset's total weight as large as possible.

Code

JavaC#

Java

C#

Complexity

Time Complexity: O(n log n) mainly from sorting and union-find operations.

Space Complexity: O(n) for the union-find structure and degree tracking.

Try this approach in the editor →

Approach 3: Default Approach

Code

PythonJavaC++GoTypeScript

Python

Java

C++

Go

TypeScript

Try this approach in the editor →

Complexity Comparison

ApproachComplexity
Approach 1: Greedy Method Using Priority Queue

Time Complexity: O(n log n) due to sorting and the find-union operations.

Space Complexity: O(n) for the disjoint-set data structure and degree tracking.

Approach 2: Kruskal's Algorithm with Modifications

Time Complexity: O(n log n) mainly from sorting and union-find operations.

Space Complexity: O(n) for the union-find structure and degree tracking.

Default Approach—

Detailed Complexity Analysis

ApproachTimeSpaceWhen to Use
Greedy DFS with Priority QueueO(n log n)O(n)Best general solution for tree constraints where child contributions must be ranked and selected.
Modified Kruskal's AlgorithmO(E log E)O(n)Useful when modeling the problem as a constrained maximum spanning structure using union–find.

Video Solution

3367. Maximize Sum of Weights after Edge Removals | DP on Trees | Memoization • Aryan Mittal • 3,293 views views

Watch 3 more video solutions →

Frequently Asked Questions

Is Maximize Sum of Weights after Edge Removals easy or hard?
Maximize Sum of Weights after Edge Removals is categorized as a Hard problem. It requires combining tree traversal, dynamic programming, and greedy selection using sorting or heaps, which makes the reasoning and implementation more complex than standard graph problems.
Maximize Sum of Weights after Edge Removals Python/Java solution
Python solutions typically use DFS with heapq to store edge gains, while Java implementations rely on PriorityQueue and adjacency lists. Both versions implement the same greedy tree DP idea and run in O(n log n) time with O(n) space.
How to solve Maximize Sum of Weights after Edge Removals in O(n log n)?
Build the tree adjacency list and run a DFS from a root. For each node, calculate the benefit of keeping each child edge versus removing it. Store these benefits in a heap and select the highest values that satisfy the allowed number of connections. Combine these selected gains with the current node’s result to propagate the optimal value upward.
What is the best approach for Maximize Sum of Weights after Edge Removals?
The most effective approach uses tree dynamic programming with a greedy selection strategy. Perform a depth‑first search and compute the gain contributed by each child edge. Store gains in a priority queue and keep only the highest‑value edges that satisfy the node constraint. This approach runs in O(n log n) time and O(n) space.
Is Maximize Sum of Weights after Edge Removals asked at Google/Amazon/Meta?
Hard tree dynamic programming and greedy selection problems like this commonly appear in interviews at companies such as Google, Amazon, and Meta. Variants involving constrained edge selection, maximum spanning structures, or tree DP frequently show up in senior‑level algorithm rounds.
What data structure is used in Maximize Sum of Weights after Edge Removals?
Key data structures include adjacency lists for representing the tree, a priority queue (heap) to rank edge gains, and recursion or stacks for DFS traversal. Some alternative solutions also use a Disjoint Set Union structure when applying a Kruskal‑style strategy.
What is the time complexity of Maximize Sum of Weights after Edge Removals?
The optimal DFS + greedy solution runs in O(n log n) time because each node may push child gains into a priority queue or perform sorting. Space complexity is O(n) for adjacency lists, recursion stack, and intermediate gain storage.

Ready to solve this problem?

Practice Maximize Sum of Weights after Edge Removals with our built-in code editor and test cases.

Practice on FleetCode
Previous

Find Minimum Diameter After Merging Two Trees