3367. Maximize Sum of Weights after Edge Removals
Description
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
kother nodes, wherekis 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 thank = 2nodes. - The sum of weights is 22, and we can't achieve a greater sum. Thus, the answer is 22.
Example 2:
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 = 3nodes, we don't remove any edges. - The sum of weights is 65. Thus, the answer is 65.
Constraints:
2 <= n <= 1051 <= k <= n - 1edges.length == n - 1edges[i].length == 30 <= edges[i][0] <= n - 10 <= edges[i][1] <= n - 11 <= edges[i][2] <= 106- The input is generated such that
edgesform a valid tree.
Solutions
Solution 1
Python3
class Solution:
def maximizeSumOfWeights(self, edges: List[List[int]], k: int) -> int:
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])
n = len(edges) + 1
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))
x, y = dfs(0, -1)
return max(x, y)
Java
class Solution {
private List<int[]>[] g;
private int k;
public long maximizeSumOfWeights(int[][] edges, int k) {
this.k = k;
int n = edges.length + 1;
g = new List[n];
Arrays.setAll(g, i -> new ArrayList<>());
for (var e : edges) {
int u = e[0], v = e[1], w = e[2];
g[u].add(new int[] {v, w});
g[v].add(new int[] {u, w});
}
var ans = dfs(0, -1);
return Math.max(ans[0], ans[1]);
}
private long[] dfs(int u, int fa) {
long s = 0;
List<Long> t = new ArrayList<>();
for (var e : g[u]) {
int v = e[0], w = e[1];
if (v == fa) {
continue;
}
var res = dfs(v, u);
s += res[0];
long d = w + res[1] - res[0];
if (d > 0) {
t.add(d);
}
}
t.sort(Comparator.reverseOrder());
for (int i = 0; i < Math.min(t.size(), k - 1); ++i) {
s += t.get(i);
}
return new long[] {s + (t.size() >= k ? t.get(k - 1) : 0), s};
}
}
C++
class Solution {
public:
long long maximizeSumOfWeights(vector<vector<int>>& edges, int k) {
int n = edges.size() + 1;
vector<vector<pair<int, int>>> g(n);
for (auto& e : edges) {
int u = e[0], v = e[1], w = e[2];
g[u].emplace_back(v, w);
g[v].emplace_back(u, w);
}
using ll = long long;
auto dfs = [&](this auto&& dfs, int u, int fa) -> pair<ll, ll> {
ll s = 0;
vector<ll> t;
for (auto& [v, w] : g[u]) {
if (v == fa) {
continue;
}
auto [a, b] = dfs(v, u);
s += a;
ll d = w + b - a;
if (d > 0) {
t.push_back(d);
}
}
ranges::sort(t, greater<>());
for (int i = 0; i < min((int) t.size(), k - 1); ++i) {
s += t[i];
}
return {s + (t.size() >= k ? t[k - 1] : 0), s};
};
auto [x, y] = dfs(0, -1);
return max(x, y);
}
};
Go
func maximizeSumOfWeights(edges [][]int, k int) int64 {
n := len(edges) + 1
g := make([][][]int, n)
for _, e := range edges {
u, v, w := e[0], e[1], e[2]
g[u] = append(g[u], []int{v, w})
g[v] = append(g[v], []int{u, w})
}
var dfs func(u, fa int) (int64, int64)
dfs = func(u, fa int) (int64, int64) {
var s int64
var t []int64
for _, e := range g[u] {
v, w := e[0], e[1]
if v == fa {
continue
}
a, b := dfs(v, u)
s += a
d := int64(w) + b - a
if d > 0 {
t = append(t, d)
}
}
sort.Slice(t, func(i, j int) bool {
return t[i] > t[j]
})
for i := 0; i < min(len(t), k-1); i++ {
s += t[i]
}
s2 := s
if len(t) >= k {
s += t[k-1]
}
return s, s2
}
x, y := dfs(0, -1)
return max(x, y)
}
TypeScript
function maximizeSumOfWeights(edges: number[][], k: number): number {
const n = edges.length + 1;
const g: [number, number][][] = Array.from({ length: n }, () => []);
for (const [u, v, w] of edges) {
g[u].push([v, w]);
g[v].push([u, w]);
}
const dfs = (u: number, fa: number): [number, number] => {
let s = 0;
const t: number[] = [];
for (const [v, w] of g[u]) {
if (v === fa) continue;
const [a, b] = dfs(v, u);
s += a;
const d = w + b - a;
if (d > 0) t.push(d);
}
t.sort((a, b) => b - a);
for (let i = 0; i < Math.min(t.length, k - 1); i++) {
s += t[i];
}
const s2 = s;
if (t.length >= k) {
s += t[k - 1];
}
return [s, s2];
};
const [x, y] = dfs(0, -1);
return Math.max(x, y);
}