Maximize sum of weights after edge removals

Dynamic Programming Tree Depth-First Search Sorting

30.5% acceptance

Feb 24, 2026

100

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.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn maximize_sum_of_weights(edges: Vec<Vec<i32>>, k: i32) -> i64 {
    let n = edges.len() + 1;
    let k = k as usize;
    
    // Build adjacency list
    let mut adj: Vec<Vec<(usize, i64)>> = vec![vec![]; n];
    for e in &edges {
      let (u, v, w) = (e[0] as usize, e[1] as usize, e[2] as i64);
      adj[u].push((v, w));
      adj[v].push((u, w));
    }
    
    // DFS post-order. For each node, we decide which child edges to keep.
    // dp(node, parent) returns (max_sum_when_edge_to_parent_is_kept, max_sum_when_edge_to_parent_not_kept)
    // For each child c with edge weight w:
    //   We can include edge(node,c): contributes w + dp_c_kept
    //   Or not include: contributes dp_c_not_kept
    // Node can keep at most k edges to children (since parent edge uses 1 slot if kept).
    
    // Returns (with_parent_edge_slot, without_parent_edge_slot)
    // with_parent: node has 1 slot used by parent, so can keep at most k-1 child edges
    // without_parent: node has no parent slot used, can keep at most k child edges
    fn dfs(node: usize, parent: usize, adj: &Vec<Vec<(usize, i64)>>, k: usize) -> (i64, i64) {
      // For each child, compute gain if we include the edge vs not
      let mut gains: Vec<i64> = Vec::new();
      let mut base = 0i64; // sum of "don't include child edge" for all children
      
      for &(c, w) in &adj[node] {
        if c == parent { continue; }
        let (c_with, c_without) = dfs(c, node, adj, k);
        // If we include edge(node,c): gain = w + c_with (child uses 1 slot for parent)
        // If we don't: c_without
        let gain = w + c_with - c_without;
        base += c_without;
        gains.push(gain);
      }
      
      // Sort gains descending; take top ones
      gains.sort_unstable_by(|a, b| b.cmp(a));
      
      // without_parent: keep top min(k, len) positive gains
      let max_no_parent = {
        let take = gains.len().min(k);
        let mut s = base;
        for i in 0..take {
          if gains[i] > 0 { s += gains[i]; } else { break; }
        }
        s
      };
      
      // with_parent: keep top min(k-1, len) positive gains (1 slot reserved for parent)
      let max_with_parent = if k >= 1 {
        let take = gains.len().min(k - 1);
        let mut s = base;
        for i in 0..take {
          if gains[i] > 0 { s += gains[i]; } else { break; }
        }
        s
      } else {
        base // k=0, no edges allowed... but k>=1 in constraints
      };
      
      (max_with_parent, max_no_parent)
    }
    
    let (_, ans) = dfs(0, usize::MAX, &adj, k);
    ans
  }
}