We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

Maximize Sum of Weights after Edge Removals

Number: 3675

Difficulty: Hard

Paid? No

Companies: N/A


Problem Description

Given an undirected tree with n nodes (numbered from 0 to n-1) and a list of weighted edges, remove zero or more edges so that for every node, the degree (number of incident edges) is at most k. The goal is to maximize the sum of weights of the remaining edges.


Key Insights

  • The tree structure implies there is no cycle so we can use DFS for dynamic programming.
  • Each node is constrained by a maximum number (k) of incident edges; if more edges exist, we must choose which ones to keep.
  • A greedy/dynamic programming approach works by, for each node, deciding which child edges yield the maximum improvement to the total sum while respecting degree limits.
  • The problem can be thought of as a kind of tree knapSack: for each node, when combining contributions from children, we choose up to k (or k-1 if an edge to the parent is used) edges that bring extra benefit relative to skipping them.
  • Using two DP states per node (one when the parent's connection is not counted and one when it is) helps manage the varying available degree at each node.

Space and Time Complexity

Time Complexity: O(n log n) in worst-case due to sorting improvements at each node (the summation over all children across nodes is O(n)). Space Complexity: O(n), due to recursion stack and the DP array for each node.


Solution

We solve the problem with a DFS-based dynamic programming approach.

  1. Represent the tree using an adjacency list.
  2. For each node, define two DP values:
    • dp[node][0] represents the maximum sum in the subtree rooted at this node when the edge from its parent is NOT considered as “using up” one of its k allowed connections (i.e. the node can use up to k edges among its children).
    • dp[node][1] represents the maximum sum in the subtree when the edge from its parent is already taken, so the node can use at most k-1 child connections.
  3. For every node, recursively calculate contributions from its children. For each child, there is a potential “gain” if we choose to include the edge (child edge weight plus dp[child][1]) over skipping the edge (dp[child][0]).
  4. Collect all positive gains from children. Then, sort them in descending order. Add as many positive gains as allowed (up to k for dp[node][0] and up to k-1 for dp[node][1]).
  5. Finally, dp[0][0] (if we take 0 as the root which has no parent) gives us the maximum sum.

Important details include ensuring that we do not count an edge twice (once for each of its two nodes) and carefully merging the gains while respecting the constraints.


Code Solutions

# Python solution with detailed comments
from collections import defaultdict
import sys
sys.setrecursionlimit(10**6)

def maximizeSumOfWeights(edges, k):
    # Build an undirected graph (tree) with edge weights.
    graph = defaultdict(list)
    for u, v, weight in edges:
        graph[u].append((v, weight))
        graph[v].append((u, weight))
    
    # DFS returns a tuple: (dp[node][0], dp[node][1])
    # dp[node][0] : maximum sum for subtree when node can choose up to k children edges.
    # dp[node][1] : maximum sum for subtree when node can choose up to k-1 children edges.
    def dfs(node, parent):
        # benefits holds the "extra" weight we get by including an edge to this child versus not including it.
        benefits = []
        total = 0
        
        for neighbor, weight in graph[node]:
            if neighbor == parent:
                continue
            # Recurse into child.
            child_dp0, child_dp1 = dfs(neighbor, node)
            # Option 1: Skip connecting to this child: use dp[child][0]
            total += child_dp0
            # Option 2: Connect this edge:
            # Gain = weight + (dp[child][1] - dp[child][0])
            gain = weight + child_dp1 - child_dp0
            benefits.append(gain)
        
        # Sort gains in descending order so that we pick the best ones.
        benefits.sort(reverse=True)
        
        # dp[node][0]: we are allowed up to k edges from this node.
        dp0 = total
        # dp[node][1]: we are allowed up to k-1 edges from this node. (Parent edge already used)
        dp1 = total
        
        # For dp0, take up to k positive gains.
        for i in range(min(k, len(benefits))):
            if benefits[i] > 0:
                dp0 += benefits[i]
            else:
                break
        
        # For dp1, take up to k-1 positive gains.
        for i in range(min(k-1, len(benefits))):
            if benefits[i] > 0:
                dp1 += benefits[i]
            else:
                break
        
        return dp0, dp1

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

# Example usage:
edges = [[0,1,4],[0,2,2],[2,3,12],[2,4,6]]
k = 2
print(maximizeSumOfWeights(edges, k))  # Expected output: 22
← Back to All Questions