Facebook Pixel

3367. Maximize Sum of Weights after Edge Removals

HardTreeDepth-First SearchDynamic Programming
Leetcode Link

Problem Description

There exists an undirected tree with n nodes numbered from 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 provided to you.
  • The sum of the weights of the remaining edges is maximized.

You need to return the maximum possible sum of weights for the remaining edges after making the necessary removals.

Intuition

The problem requires us to ensure that no node in the tree has more than k connections while maximizing the sum of the remaining edge weights. To achieve this, we can use a depth-first search (DFS) approach combined with a greedy strategy.

The idea is to calculate the maximum sum of weights that can be retained by making sure that any node connects at most k other nodes:

  1. Tree Traversal: Start by performing DFS from any node, typically the root (let's select node 0).

  2. Recalculating the Sum and Differences: As we traverse, compute two values via DFS:

    • s is a running sum of the weights that are retained.
    • t is a list that stores the difference between the weight of an edge and the potential gain if that edge were omitted.

  3. Greedy Selection: For each node, sort the potential gain from omitting edges. Use the top k (where applicable without exceeding the bounds of t) to maximize the sum of remaining edge weights based on possible connectivity.

  4. Result Computation: Each DFS call returns two possible sums calculated based on keeping or not keeping certain edges. The maximum possible between the two ensures the optimal solution is considered.

By strategically deciding which edges to retain during the DFS traversal by considering local gains, we maximize the total weight of the resulting tree while respecting the constraints.

Learn more about Tree, Depth-First Search and Dynamic Programming patterns.

Solution Approach

The solution to the problem leverages a combination of DFS (Depth-First Search) traversal and a greedy selection strategy, to ensure that each node adheres to the constraint of connecting to no more than k other nodes while maximizing the sum of the edge weights.

Here's a breakdown of the implementation approach:

  1. Graph Representation:

    • We represent the tree using an adjacency list. This is a list of lists g where g[i] contains tuples (v, w) for an edge between node i and node v with weight w.
    g: List[List[Tuple[int, int]]] = [[] for _ in range(n)]
for u, v, w in edges:
    g[u].append((v, w))
    g[v].append((u, w))

  2. DFS Function:

    • We define a dfs function that performs a traversal starting from a node u, avoiding the parent node fa. This function returns two sums: one including and one excluding specific adjustments.
    def dfs(u: int, fa: int) -> Tuple[int, int]:
    s = 0
    t = []
    for v, w in g[u]:
        if v == fa:
            continue
        a, b = dfs(v, u)
        s += a
        if (d := (w + b - a)) > 0:
            t.append(d)
    t.sort(reverse=True)
    return s + sum(t[:k]), s + sum(t[: k - 1])

  3. Sum Calculation:

    • For each node, compute the sum s of weights for its subtree.
    • Compute list t, storing potential gains from keeping certain edges.
    • Sort t in descending order to select the top values that maximize the sum within the constraint of at most k connected nodes.

  4. Result Extraction:

    • After running dfs starting from the root node (usually node 0), we obtain two possible sums: x (maximized sum including any top k differences) and y (maximized for one less than k connections).
    • The final result is the maximum of these two, ensuring optimal selection based on the constraint.
    x, y = dfs(0, -1)
return max(x, y)

This approach carefully balances the need to maximize edge weights while respecting the connectivity limit imposed by k using a combination of recursion, greedy selection, and depth-first search.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a small example to illustrate the solution approach:

Consider a tree with 4 nodes and the following edges:
edges = [[0, 1, 5], [1, 2, 4], [1, 3, 3]]
We have n = 4 nodes and we want to ensure each node connects to at most k = 2 other nodes.

  1. Graph Representation:

    • Start by representing the tree as an adjacency list:
    g = [
  [(1, 5)],       # Node 0 connects to Node 1 with weight 5.
  [(0, 5), (2, 4), (3, 3)],  # Node 1 connects to Nodes 0, 2, 3 with weights 5, 4, 3.
  [(1, 4)],       # Node 2 connects to Node 1 with weight 4.
  [(1, 3)]        # Node 3 connects to Node 1 with weight 3.
]

  2. DFS Traversal:

    • Begin DFS from the root node (0). We ignore the parent node to avoid cycles.
    • At each node, we evaluate keeping or omitting edges based on maximizing weight sum.

  3. Node Analysis:

    • Start DFS at node 0:

      • Node 0 connects to Node 1 with weight 5.
      • Recursive DFS call on Node 1 excluding Node 0.

    • At Node 1, it connects to Nodes 0, 2, 3 (already visited 0):

      • Recursive DFS calls on Node 2 and Node 3.

    • For Node 2:

      • Connects back to Node 1 with weight 4.
      • Only one connection; return (0, 0).

    • For Node 3:

      • Connects back to Node 1 with weight 3.
      • Only one connection; return (0, 0).

  4. Edge Evaluation and Selection:

    • Returning to Node 1 from its children with:
      • Weights: 4 from Node 2, 3 from Node 3.
      • Sort edges based on potential gain if they are omitted: [(4, 0), (3, 0)].
      • Keep top k=2 connections by selecting weights 4 and 3.

  5. Result Calculation:

    • Backtrack to Node 0, aggregate weight:
      • Include Node 1 -> 2 and Node 1 -> 3 with weights 4 and 3.
      • Final weight if edge 0 -> 1 is retained: 5 + 4 + 3 = 12.

The maximum possible weight sum while ensuring each node has at most 2 connections is 12.

Solution Implementation

1from typing import List, Tuple
2
3class Solution:
4    def maximizeSumOfWeights(self, edges: List[List[int]], k: int) -> int:
5        def dfs(node: int, parent: int) -> Tuple[int, int]:
6            current_sum = 0  # Holds the sum of selected weights
7            differences = []  # Stores the beneficial weight differences
8          
9            # Traverse all connected nodes (children) from the current node
10            for neighbor, weight in graph[node]:
11                if neighbor == parent:
12                    continue  # Skip the parent node to avoid cycles
13              
14                subtree_sum, subtree_alt_sum = dfs(neighbor, node)
15                current_sum += subtree_sum
16              
17                # Calculate the weight difference if adding this edge is beneficial
18                difference = weight + subtree_alt_sum - subtree_sum
19                if difference > 0:
20                    differences.append(difference)
21          
22            # Sort differences in descending order to maximize the sum
23            differences.sort(reverse=True)
24          
25            # Calculate the maximum possible sum by selecting top k differences
26            max_sum_with_k = current_sum + sum(differences[:k])
27            max_sum_with_k_minus_one = current_sum + sum(differences[:k - 1])
28          
29            return max_sum_with_k, max_sum_with_k_minus_one
30
31        # Number of nodes is 1 more than the number of edges for a tree
32        num_nodes = len(edges) + 1
33      
34        # Initialize the graph as an adjacency list
35        graph: List[List[Tuple[int, int]]] = [[] for _ in range(num_nodes)]
36      
37        # Build the graph from the given edges
38        for u, v, weight in edges:
39            graph[u].append((v, weight))
40            graph[v].append((u, weight))
41      
42        # Start DFS from the root node, which is considered node 0
43        max_weight, alt_weight = dfs(0, -1)
44      
45        # Return the best possible sum
46        return max(max_weight, alt_weight)
47
1import java.util.ArrayList;
2import java.util.Arrays;
3import java.util.Comparator;
4import java.util.List;
5
6class Solution {
7    private List<int[]>[] graph; // adjacency list to store graph
8    private int k; // number of extra edges allowed
9
10    public long maximizeSumOfWeights(int[][] edges, int k) {
11        this.k = k;
12        int n = edges.length + 1; // number of nodes
13        graph = new List[n];
14      
15        // Initialize the adjacency list
16        Arrays.setAll(graph, i -> new ArrayList<>());
17      
18        // Build the graph from edge list
19        for (var edge : edges) {
20            int u = edge[0], v = edge[1], w = edge[2]; // u and v are nodes, w is the weight of the edge
21            graph[u].add(new int[]{v, w});
22            graph[v].add(new int[]{u, w});
23        }
24      
25        // Start depth-first search from node 0 assuming it's the root
26        var result = dfs(0, -1);
27      
28        // Return the maximum of including or not including the k-th maximum extra weight edge
29        return Math.max(result[0], result[1]);
30    }
31
32    private long[] dfs(int currentNode, int parent) {
33        long sum = 0; // sum of the maximum weights for the current subtree
34        List<Long> extraWeights = new ArrayList<>(); // potential extra weights from different paths
35
36        // Traverse all adjacent nodes
37        for (var edge : graph[currentNode]) {
38            int adjacentNode = edge[0], weight = edge[1];
39          
40            // Skip the parent to avoid going back in the DFS
41            if (adjacentNode == parent) {
42                continue;
43            }
44          
45            // Recursively process the adjacent node's subtree
46            var subtreeResult = dfs(adjacentNode, currentNode);
47
48            // Add the best weight of the subtree rooted at 'adjacentNode' to the current sum
49            sum += subtreeResult[0];
50          
51            // Calculate the difference between choosing and not choosing this path
52            long delta = weight + subtreeResult[1] - subtreeResult[0];
53            if (delta > 0) {
54                extraWeights.add(delta);
55            }
56        }
57
58        // Sort potential extra weights in descending order
59        extraWeights.sort(Comparator.reverseOrder());
60      
61        // Pick the top (k-1) extra weights to maximize the sum
62        for (int i = 0; i < Math.min(extraWeights.size(), k - 1); ++i) {
63            sum += extraWeights.get(i);
64        }
65      
66        // Return a pair [maximum sum including k extra paths, sum without including the k-th path]
67        return new long[]{sum + (extraWeights.size() >= k ? extraWeights.get(k - 1) : 0), sum};
68    }
69}
70
1#include <vector>
2#include <algorithm>
3#include <functional>
4
5class Solution {
6public:
7    long long maximizeSumOfWeights(std::vector<std::vector<int>>& edges, int k) {
8        int n = edges.size() + 1;  // Number of nodes in the graph
9        std::vector<std::vector<std::pair<int, int>>> graph(n); // Adjacency list representation
10
11        // Construct the graph from the edges
12        for (auto& edge : edges) {
13            int u = edge[0], v = edge[1], weight = edge[2];
14            graph[u].emplace_back(v, weight);
15            graph[v].emplace_back(u, weight);
16        }
17
18        using ll = long long;
19
20        // Define the DFS function to calculate the maximum sum of weights
21        std::function<std::pair<ll, ll>(int, int)> dfs = [&](int node, int parent) -> std::pair<ll, ll> {
22            ll subtree_sum = 0;
23            std::vector<ll> differences;
24
25            // Explore all adjacent nodes
26            for (auto& [adj_node, weight] : graph[node]) {
27                if (adj_node == parent) {
28                    continue;  // Skip the parent node
29                }
30                auto [a, b] = dfs(adj_node, node);  // Recursive DFS call
31                subtree_sum += a;
32                ll diff = weight + b - a;  // Calculate weight difference
33                if (diff > 0) {
34                    differences.push_back(diff);
35                }
36            }
37
38            // Sort differences in descending order
39            std::sort(differences.begin(), differences.end(), std::greater<ll>());
40
41            // Select up to k-1 differences to maximize the sum
42            for (int i = 0; i < std::min(static_cast<int>(differences.size()), k - 1); ++i) {
43                subtree_sum += differences[i];
44            }
45
46            ll max_difference = (differences.size() >= k) ? differences[k - 1] : 0;  // Select the k-th largest difference if it exists
47            return {subtree_sum + max_difference, subtree_sum};
48        };
49
50        auto [max_sum_with_root, max_sum_without_root] = dfs(0, -1);  // Start DFS from node 0
51        return std::max(max_sum_with_root, max_sum_without_root);  // Return the maximum of two possible sums
52    }
53};
54
1/**
2 * Computes the maximum sum of weights after selecting edges.
3 * @param edges A 2D array where each element is an array consisting of [u, v, w] indicating an edge from u to v with weight w.
4 * @param k The maximum number of edges you can choose.
5 * @returns The maximum sum of selected edge weights.
6 */
7function maximizeSumOfWeights(edges: number[][], k: number): number {
8    // n is the number of nodes in the graph
9    const n = edges.length + 1;
10    // graph represented as an adjacency list, each node has a list of tuples [neighbor, weight]
11    const adjacencyList: [number, number][][] = Array.from({ length: n }, () => []);
12
13    // Build the graph from the edges
14    for (const [u, v, weight] of edges) {
15        adjacencyList[u].push([v, weight]);
16        adjacencyList[v].push([u, weight]);
17    }
18
19    /**
20     * Depth First Search to calculate the maximum sum of weights
21     * @param node The current node.
22     * @param parent The parent node to avoid revisiting.
23     * @returns A tuple containing two values: the maximum sum including k edges and the sum excluding the last potential kth edge.
24     */
25    const dfs = (node: number, parent: number): [number, number] => {
26        let sum = 0; // Tracks the maximal sum
27        const potentialIncreases: number[] = []; // Potential increases for the sum when adding one more edge
28
29        // Explore neighbors of the current node
30        for (const [neighbor, weight] of adjacencyList[node]) {
31            // Skip the parent node to prevent revisiting
32            if (neighbor === parent) continue;
33
34            // Perform DFS on the child node
35            const [maxSumIncluding, maxSumExcluding] = dfs(neighbor, node);
36
37            // Accumulate the maximum sum considering current path
38            sum += maxSumIncluding;
39
40            // Calculate possible increase and consider if beneficial
41            const increase = weight + maxSumExcluding - maxSumIncluding;
42            if (increase > 0) potentialIncreases.push(increase);
43        }
44
45        // Sort potential increases in descending order to maximize benefit
46        potentialIncreases.sort((a, b) => b - a);
47
48        // Include the best (k-1) edges' increases to maximize the total sum
49        for (let i = 0; i < Math.min(potentialIncreases.length, k - 1); i++) {
50            sum += potentialIncreases[i];
51        }
52      
53        // Store the total sum when k-1 edges are considered
54        const sumExcludingLastEdge = sum;
55
56        // If there's a kth edge to consider, include its possible increase
57        if (potentialIncreases.length >= k) {
58            sum += potentialIncreases[k - 1];
59        }
60
61        // Return both sum with and without the last possible edge
62        return [sum, sumExcludingLastEdge];
63    };
64
65    // Start DFS from node 0 with no parent (-1 indicates no parent)
66    const [maxSumWithK, maxSumWithoutLast] = dfs(0, -1);
67
68    // Return the maximum of the two calculated sums
69    return Math.max(maxSumWithK, maxSumWithoutLast);
70}
71

Time and Space Complexity

The time complexity of the code is O(n log n), where n is the number of nodes in the tree. This is because the DFS traversal visits each node once, and sorting the list t has a complexity of O(n log n) in the worst case when all potential differences d are stored in t.

The space complexity is O(n). This arises from the adjacency list g, which stores the tree structure and occupies O(n) space, and the recursion stack used by the dfs function, which also takes up O(n) space in the worst case when the tree is skewed.

Learn more about how to find time and space complexity quickly.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

Which of the two traversal algorithms (BFS and DFS) can be used to find whether two nodes are connected?

Recommended Readings

Everything About Trees

Everything About Trees A tree is a type of graph data structure composed of nodes and edges Its main properties are It is acyclic doesn't contain any cycles There exists a path from the root to any node Has N 1 edges where N is the number of nodes in the tree and

DFS Fundamentals

https algomonster s3 us east 2 amazonaws com cover_photos dfs svg Depth First Search Prereqs Recursion Review problems recursion_intro Trees problems tree_intro With a solid understanding of recursion under our belts we are now ready to tackle one of the most useful techniques in coding interviews Depth First Search DFS

Dynamic Programming Introduction and Patterns

What is Dynamic Programming Prerequisite DFS problems dfs_intro Backtracking problems backtracking Memoization problems memoization_intro Pruning problems backtracking_pruning Dynamic programming is an algorithmic optimization technique that breaks down a complicated problem into smaller overlapping sub problems in a recursive manner and uses solutions to the sub problems to construct a solution

Want a Structured Path to Master System Design Too? Don’t Miss This!

Load More