3367. Maximize Sum of Weights after Edge Removals LeetCode Solution
In this guide, you will get 3367. Maximize Sum of Weights after Edge Removals LeetCode Solution with the best time and space complexity. The solution to Maximize Sum of Weights after Edge Removals problem is provided in various programming languages like C++, Java, and Python. This will be helpful for you if you are preparing for placements, hackathons, interviews, or practice purposes. The solutions provided here are very easy to follow and include detailed explanations.
Problem Statement of Maximize Sum of Weights after Edge Removals
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.
Example 1:
Input: edges = [[0,1,4],[0,2,2],[2,3,12],[2,4,6]], k = 2
Output: 22
Explanation:
Node 2 has edges with 3 other nodes. We remove the edge [0, 2, 2], ensuring that no node has edges with more than k = 2 nodes.
The sum of weights is 22, and we can’t achieve a greater sum. Thus, the answer is 22.
Input: edges = [[0,1,5],[1,2,10],[0,3,15],[3,4,20],[3,5,5],[0,6,10]], k = 3
Output: 65
Explanation:
Since no node has edges connecting it to more than k = 3 nodes, we don’t remove any edges.
The sum of weights is 65. Thus, the answer is 65.
Constraints:
2 <= n <= 105
1 <= k <= n – 1
edges.length == n – 1
edges[i].length == 3
0 <= edges[i][0] <= n – 1
0 <= edges[i][1] <= n – 1
1 <= edges[i][2] <= 106
The input is generated such that edges form a valid tree.
Complexity Analysis
Time Complexity: O(n\log k)
Space Complexity: O(n)
3367. Maximize Sum of Weights after Edge Removals LeetCode Solution in C++
class Solution {
public:
long long maximizeSumOfWeights(vector<vector<int>>& edges, int k) {
const int n = edges.size() + 1;
vector<vector<pair<int, int>>> graph(n);
for (const vector<int>& edge : edges) {
const int u = edge[0];
const int v = edge[1];
const int w = edge[2];
graph[u].emplace_back(v, w);
graph[v].emplace_back(u, w);
}
return dfs(graph, 0, -1, k).second;
}
// Returns
// (the weight sum of the subtree rooted at u with at most k - 1 children,
// the weight sum of the subtree rooted at u with at most k children).
pair<long, long> dfs(const vector<vector<pair<int, int>>>& graph, int u,
int prev, int k) {
long weightSum = 0;
priority_queue<long> diffs;
for (const auto& [v, w] : graph[u]) {
if (v == prev)
continue;
const auto [subK1, subK] = dfs(graph, v, u, k);
weightSum += subK;
// If picking (u, v) makes the sum larger, we should pick it.
diffs.push(max(0L, subK1 - subK + w));
}
long topK1 = 0;
long topK = 0;
for (int i = 0; i < k && !diffs.empty(); ++i) {
if (i < k - 1)
topK1 += diffs.top();
topK += diffs.top();
diffs.pop();
}
return {weightSum + topK1, weightSum + topK};
};
}; /* code provided by PROGIEZ */
3367. Maximize Sum of Weights after Edge Removals LeetCode Solution in Java
class Solution {
public long maximizeSumOfWeights(int[][] edges, int k) {
final int n = edges.length + 1;
List<Pair<Integer, Integer>>[] graph = new List[n];
for (int i = 0; i < n; ++i)
graph[i] = new ArrayList<>();
for (int[] edge : edges) {
final int u = edge[0];
final int v = edge[1];
final int w = edge[2];
graph[u].add(new Pair<>(v, w));
graph[v].add(new Pair<>(u, w));
}
return dfs(graph, 0, -1, k).getValue();
}
// Returns
// (the weight sum of the subtree rooted at u with at most k - 1 children,
// the weight sum of the subtree rooted at u with at most k children).
private Pair<Long, Long> dfs(List<Pair<Integer, Integer>>[] graph, int u, int prev, int k) {
long weightSum = 0;
Queue<Long> diffs = new PriorityQueue<>(Collections.reverseOrder());
for (Pair<Integer, Integer> pair : graph[u]) {
final int v = pair.getKey();
final int w = pair.getValue();
if (v == prev)
continue;
Pair<Long, Long> subResult = dfs(graph, v, u, k);
final long subK1 = subResult.getKey();
final long subK = subResult.getValue();
weightSum += subK;
// If picking (u, v) makes the sum larger, we should pick it.
diffs.offer(Math.max(0L, subK1 - subK + w));
}
long topK1 = 0;
long topK = 0;
for (int i = 0; i < k && !diffs.isEmpty(); ++i) {
if (i < k - 1)
topK1 += diffs.peek();
topK += diffs.poll();
}
return new Pair<>(weightSum + topK1, weightSum + topK);
}
} // code provided by PROGIEZ
3367. Maximize Sum of Weights after Edge Removals LeetCode Solution in Python
class Solution:
def maximizeSumOfWeights(self, edges: list[list[int]], k: int) -> int:
graph = [[] for _ in range(len(edges) + 1)]
for u, v, w in edges:
graph[u].append((v, w))
graph[v].append((u, w))
def dfs(u: int, prev: int) -> tuple[int, int]:
"""
Returns
(the weight sum of the subtree rooted at u with at most k - 1 children,
the weight sum of the subtree rooted at u with at most k children).
"""
weightSum = 0
diffs = []
for v, w in graph[u]:
if v == prev:
continue
subK1, subK = dfs(v, u)
weightSum += subK
# If picking (u, v) makes the sum larger, we should pick it.
diffs.append(max(0, subK1 - subK + w))
return (weightSum + sum(heapq.nlargest(k - 1, diffs)),
weightSum + sum(heapq.nlargest(k, diffs)))
return dfs(0, -1)[1] # code by PROGIEZ
