Maximize Sum of Weights after Edge Removals

Maximize Sum of Weights after Edge Removals - Problem

Maximize Sum of Weights after Edge Removals

You are given an undirected tree with n nodes numbered from 0 to n - 1. The tree is represented by a 2D array edges where edges[i] = [u_i, v_i, w_i] indicates there is an edge between nodes u_i and v_i with weight w_i.

Your goal is to strategically remove some edges such that:
• Each remaining node has at most k connections
• The total weight of remaining edges is maximized

Think of it as pruning a weighted tree to satisfy degree constraints while keeping the most valuable connections. This is a classic tree dynamic programming problem that requires careful consideration of each node's optimal edge selection.

Return the maximum possible sum of weights for the remaining edges.

Input & Output

example_1.py — Basic Tree

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

› Output: 22

💡 Note: We can keep edges (2,3) with weight 12, (3,4) with weight 6, and (0,1) with weight 4. Each node has at most 2 connections and total weight is 12+6+4=22.

example_2.py — Star Graph

$ Input: edges = [[0,1,10],[0,2,20],[0,3,15],[0,4,5]], k = 2

› Output: 35

💡 Note: Node 0 can keep at most 2 edges. We choose the heaviest ones: (0,2) with weight 20 and (0,3) with weight 15. Total weight is 20+15=35.

example_3.py — Small Tree

$ Input: edges = [[0,1,1],[1,2,3]], k = 1

› Output: 3

💡 Note: With k=1, we can keep at most one edge per node. The optimal choice is edge (1,2) with weight 3, giving us maximum weight of 3.

Constraints

1 ≤ n ≤ 10⁵
1 ≤ k ≤ n-1
edges.length == n-1
0 ≤ u_i, v_i < n
1 ≤ w_i ≤ 10⁶
The input forms a valid tree (connected and acyclic)

Visualization

Tap to expand

Understanding the Visualization

Model as Tree

Each server is a node, each network cable is a weighted edge

Apply Constraints

Each server can maintain at most k active connections due to bandwidth limits

Greedy Selection

For each server, choose the k highest-bandwidth connections to maximize total capacity

DFS Coordination

Use post-order traversal to ensure child decisions inform parent choices

Key Takeaway

🎯 Key Insight: Tree DP with post-order DFS allows each node to optimally select its k best edges while coordinating with the overall tree structure

Asked in

G Google 42 a Amazon 38 f Meta 25 ⊞ Microsoft 31

The optimal solution uses Tree Dynamic Programming with DFS traversal. Key insight: perform post-order traversal and at each node, greedily select the k heaviest edges from available neighbors. This achieves O(n log n) time complexity by sorting edges at each node.

Common Approaches

Approach	Time	Space	Notes
✓ Tree DP with DFS (Optimal)	O(n log n)	O(n)	Use dynamic programming on the tree structure to optimally select edges

Tree DP with DFS (Optimal) — Algorithm Steps

Build adjacency list representation of the tree
Perform post-order DFS starting from any node as root
For each node, collect all potential edges (to parent and children)
Greedily select the top k edges by weight for each node
Sum up the weights, being careful not to double-count edges

Visualization

Tap to expand

Step-by-Step Walkthrough

Build Graph

Create adjacency list representation of the tree

DFS Traversal

Perform post-order traversal, processing children before parent

Edge Selection

At each node, greedily select the k highest-weight edges

Accumulate

Sum up selected edge weights to get final answer

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>

typedef struct {
    int neighbor;
    int weight;
} Edge;

typedef struct {
    Edge* edges;
    int count;
    int capacity;
} AdjList;

AdjList* graph;
bool* visited;
long long totalWeight = 0;

int compareEdges(const void* a, const void* b) {
    Edge* ea = (Edge*)a;
    Edge* eb = (Edge*)b;
    return eb->weight - ea->weight; // Descending order
}

void dfs(int node, int k, int n) {
    visited[node] = true;
    
    // Collect unvisited neighbors
    Edge* candidates = malloc(graph[node].count * sizeof(Edge));
    int candidateCount = 0;
    
    for (int i = 0; i < graph[node].count; i++) {
        if (!visited[graph[node].edges[i].neighbor]) {
            candidates[candidateCount++] = graph[node].edges[i];
        }
    }
    
    // Sort by weight descending
    qsort(candidates, candidateCount, sizeof(Edge), compareEdges);
    
    // Take at most k edges
    int selected = (k < candidateCount) ? k : candidateCount;
    for (int i = 0; i < selected; i++) {
        totalWeight += candidates[i].weight;
        dfs(candidates[i].neighbor, k, n);
    }
    
    free(candidates);
}

long long maximumWeight(int edges[][3], int edgeCount, int k) {
    if (edgeCount == 0) return 0;
    
    int n = edgeCount + 1;
    graph = malloc(n * sizeof(AdjList));
    visited = calloc(n, sizeof(bool));
    totalWeight = 0;
    
    // Initialize adjacency list
    for (int i = 0; i < n; i++) {
        graph[i].edges = malloc(n * sizeof(Edge));
        graph[i].count = 0;
        graph[i].capacity = n;
    }
    
    // Build adjacency list
    for (int i = 0; i < edgeCount; i++) {
        int u = edges[i][0], v = edges[i][1], w = edges[i][2];
        graph[u].edges[graph[u].count++] = (Edge){v, w};
        graph[v].edges[graph[v].count++] = (Edge){u, w};
    }
    
    dfs(0, k, n);
    
    // Cleanup
    for (int i = 0; i < n; i++) {
        free(graph[i].edges);
    }
    free(graph);
    free(visited);
    
    return totalWeight;
}

int main() {
    printf("Tree DP implementation\n");
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n log n)

O(n) for DFS traversal, O(log n) for sorting edges at each node

⚡ Linearithmic

Space Complexity

O(n)

Recursion stack and adjacency list storage

⚡ Linearithmic Space

23.4K Views

Medium-High Frequency

~25 min Avg. Time

892 Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

Tree DP with DFS (Optimal) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler