-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathSolution.cs
More file actions
43 lines (42 loc) · 1.23 KB
/
Solution.cs
File metadata and controls
43 lines (42 loc) · 1.23 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
public class Solution
{
public long MaximizeSumOfWeights(int[][] edges, int k)
{
Dictionary<int, List<(int v, int w)>> graph = [];
foreach (var edge in edges)
{
int u = edge[0], v = edge[1], w = edge[2];
if (!graph.ContainsKey(u))
{
graph[u] = [];
}
graph[u].Add((v, w));
if (!graph.ContainsKey(v))
{
graph[v] = [];
}
graph[v].Add((u, w));
}
return DFS(graph, 0, -1, k)[0];
}
private long[] DFS(Dictionary<int, List<(int v, int w)>> graph, int u, int pre, int k)
{
long sum = 0;
PriorityQueue<long, long> pq = new();
foreach (var (v, w) in graph[u])
{
if (v != pre)
{
long[] next = DFS(graph, v, u, k);
next[1] += w;
sum += Math.Max(next[0], next[1]);
if (next[0] < next[1])
{
pq.Enqueue(next[1] - next[0], next[1] - next[0]);
sum -= pq.Count > k ? pq.Dequeue() : 0;
}
}
}
return [sum, sum - (pq.Count < k ? 0 : pq.Dequeue())];
}
}