Maximize Sum of Weights after Edge Removals - Problem

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] = [u_i, v_i, w_i] indicates that there is an edge between nodes u_i and v_i with weight w_i 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.

Input & Output

Example 1 — Basic Tree with k=2

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

› Output: 7

💡 Note: Node 1 has degree 3 but k=2, so we must remove at least 1 edge. To maximize weight, keep the two highest weight edges from node 1: [0,1,4] and [1,3,3] for total weight 4+3=7.

Example 2 — Degree Constraint Forces Removal

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

› Output: 7

💡 Note: Node 1 connects to 4 other nodes but can only keep 2 connections. To maximize weight, keep the two highest weight edges: [1,4,4] and [1,3,3] for total weight 4+3=7.

Example 3 — Single Edge

$ Input: edges = [[0,1,10]], k = 1

› Output: 10

💡 Note: Only one edge exists and both nodes have degree 1 ≤ k=1, so keep the edge with weight 10.

Constraints

1 ≤ n ≤ 10⁴
0 ≤ edges.length ≤ min(n × (n - 1) / 2, 10⁴)
edges[i].length == 3
0 ≤ u_i < v_i < n
1 ≤ w_i ≤ 10⁶
1 ≤ k ≤ n - 1

Visualization

Tap to expand

Asked in

G Google 15 M Microsoft 12 A Amazon 8

The optimal approach uses tree DP with DFS traversal to select the best k edges for each node while considering parent connections. The key insight is that for each node, we can keep at most k edges total, and if we connect to our parent, we can only keep k-1 child edges. Time complexity: O(n log n), Space complexity: O(n).

Common Approaches

✓ Binary Search Answer

⏱️ Time: N/A Space: N/A

Brute Force - Try All Combinations

⏱️ Time: O(2^n × n) Space: O(n)

Try every possible combination of edges to keep, check if degree constraints are satisfied, and return the maximum weight among valid combinations.

Greedy Edge Selection

⏱️ Time: O(n log n) Space: O(n)

Sort edges by weight in descending order and greedily add edges to the result as long as both endpoints have degree less than k.

Tree DP with DFS

⏱️ Time: O(n log n) Space: O(n)

Perform DFS on the tree and use dynamic programming to select the best k edges for each node, considering whether it connects to its parent or not.

Algorithm Steps — Algorithm Steps

Code -

solution.c — C

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

#define MAX_NODES 1000
#define MAX_EDGES 10000

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

typedef struct {
    Edge edges[MAX_EDGES];
    int count;
} AdjList;

static AdjList graph[MAX_NODES];
static char kept_edges[MAX_EDGES * 10][100];
static int kept_count = 0;
static char final_edges[MAX_EDGES * 10][100];
static int final_count = 0;

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

bool containsEdge(char edges[][100], int count, char* edge) {
    for (int i = 0; i < count; i++) {
        if (strcmp(edges[i], edge) == 0) {
            return true;
        }
    }
    return false;
}

void addEdge(char edges[][100], int* count, char* edge) {
    if (!containsEdge(edges, *count, edge)) {
        strcpy(edges[*count], edge);
        (*count)++;
    }
}

int min(int a, int b) {
    return a < b ? a : b;
}

int max(int a, int b) {
    return a > b ? a : b;
}

int solution(int edges[][3], int edgeCount, int k) {
    // Initialize graph
    for (int i = 0; i < MAX_NODES; i++) {
        graph[i].count = 0;
    }
    kept_count = 0;
    final_count = 0;
    
    // Build adjacency list with weights
    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};
    }
    
    // For each node, keep only the k highest-weight edges
    for (int node = 0; node < MAX_NODES; node++) {
        if (graph[node].count == 0) continue;
        
        // Sort by weight in descending order
        qsort(graph[node].edges, graph[node].count, sizeof(Edge), compareEdges);
        
        int limit = min(k, graph[node].count);
        for (int i = 0; i < limit; i++) {
            int neighbor = graph[node].edges[i].neighbor;
            int weight = graph[node].edges[i].weight;
            // Add edge in canonical form (smaller node first)
            char edge[100];
            sprintf(edge, "%d,%d,%d", min(node, neighbor), max(node, neighbor), weight);
            addEdge(kept_edges, &kept_count, edge);
        }
    }
    
    // Now we need to check which edges are kept by both endpoints
    for (int i = 0; i < edgeCount; i++) {
        int u = edges[i][0], v = edges[i][1], w = edges[i][2];
        // Check if this edge is kept by both u and v
        char edge_canonical[100];
        sprintf(edge_canonical, "%d,%d,%d", min(u, v), max(u, v), w);
        
        if (containsEdge(kept_edges, kept_count, edge_canonical)) {
            // Check if both u and v want to keep this edge
            Edge u_edges[MAX_EDGES];
            int u_count = graph[u].count;
            for (int j = 0; j < u_count; j++) {
                u_edges[j] = graph[u].edges[j];
            }
            qsort(u_edges, u_count, sizeof(Edge), compareEdges);
            int u_limit = min(k, u_count);
            
            Edge v_edges[MAX_EDGES];
            int v_count = graph[v].count;
            for (int j = 0; j < v_count; j++) {
                v_edges[j] = graph[v].edges[j];
            }
            qsort(v_edges, v_count, sizeof(Edge), compareEdges);
            int v_limit = min(k, v_count);
            
            bool u_wants = false;
            for (int j = 0; j < u_limit; j++) {
                if (u_edges[j].neighbor == v && u_edges[j].weight == w) {
                    u_wants = true;
                    break;
                }
            }
            
            bool v_wants = false;
            for (int j = 0; j < v_limit; j++) {
                if (v_edges[j].neighbor == u && v_edges[j].weight == w) {
                    v_wants = true;
                    break;
                }
            }
            
            if (u_wants && v_wants) {
                addEdge(final_edges, &final_count, edge_canonical);
            }
        }
    }
    
    // Calculate total weight
    int total_weight = 0;
    for (int i = 0; i < final_count; i++) {
        char* token = strtok(final_edges[i], ",");
        token = strtok(NULL, ",");
        token = strtok(NULL, ",");
        int w = atoi(token);
        total_weight += w;
        // Restore the string for potential future use
        sprintf(final_edges[i], "%s", final_edges[i]);
    }
    return total_weight;
}

void parseArray(const char* str, int arr[][3], int* size) {
    *size = 0;
    const char* p = str;
    while (*p && *p != '[') p++;
    if (*p == '[') p++;
    
    while (*p && *p != ']') {
        while (*p == ' ' || *p == ',') p++;
        if (*p == '[') {
            p++;
            for (int j = 0; j < 3; j++) {
                while (*p == ' ') p++;
                arr[*size][j] = (int)strtol(p, (char**)&p, 10);
                while (*p == ' ' || *p == ',') p++;
            }
            while (*p == ']') p++;
            (*size)++;
        } else {
            p++;
        }
    }
}

int main() {
    char edgesStr[10000];
    fgets(edgesStr, sizeof(edgesStr), stdin);
    edgesStr[strcspn(edgesStr, "\n")] = 0;
    
    int k;
    scanf("%d", &k);
    
    int edges[MAX_EDGES][3];
    int edgeCount;
    parseArray(edgesStr, edges, &edgeCount);
    
    int result = solution(edges, edgeCount, k);
    printf("%d\n", result);
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

⚡ Linearithmic

Space Complexity

⚡ Linearithmic Space

8.5K Views

Medium Frequency

~35 min Avg. Time

248 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