Improved task 3367

Author: javadev

src/main/kotlin/g3301_3400/s3367_maximize_sum_of_weights_after_edge_removals/Solution.kt

@@ -1,41 +1,50 @@
 package g3301_3400.s3367_maximize_sum_of_weights_after_edge_removals
 
-// #Hard #2024_11_24_Time_220_ms_(100.00%)_Space_154.5_MB_(100.00%)
+// #Hard #2024_11_27_Time_118_ms_(100.00%)_Space_140.5_MB_(100.00%)
 
 import java.util.PriorityQueue
 import kotlin.math.max
 
 class Solution {
+    private lateinit var adj: Array<MutableList<IntArray>>
+    private var k = 0
+
     fun maximizeSumOfWeights(edges: Array<IntArray>, k: Int): Long {
-        val map = HashMap<Int?, ArrayList<IntArray>?>()
-        for (edge in edges) {
-            map.computeIfAbsent(edge[0]) { _: Int? -> ArrayList<IntArray>() }!!
-                .add(intArrayOf(edge[1], edge[2]))
-            map.computeIfAbsent(edge[1]) { _: Int? -> ArrayList<IntArray>() }!!
-                .add(intArrayOf(edge[0], edge[2]))
+        val n = edges.size + 1
+        adj = Array(n) { ArrayList<IntArray>() }
+        this.k = k
+        for (i in 0..<n) {
+            adj[i] = ArrayList<IntArray>()
+        }
+        for (e in edges) {
+            adj[e[0]].add(e)
+            adj[e[1]].add(e)
         }
-        return maximizeSumOfWeights(0, -1, k, map)[0]
+        return dfs(0, -1)[1]
     }
 
-    private fun maximizeSumOfWeights(
-        v: Int,
-        from: Int,
-        k: Int,
-        map: HashMap<Int?, ArrayList<IntArray>?>,
-    ): LongArray {
+    private fun dfs(v: Int, parent: Int): LongArray {
         var sum: Long = 0
-        val queue = PriorityQueue<Long>()
-        for (i in map[v]!!) {
-            if (i[0] != from) {
-                val next = maximizeSumOfWeights(i[0], v, k, map)
-                next[1] += i[1].toLong()
-                sum = sum + max(next[0], next[1])
-                if (next[0] < next[1]) {
-                    queue.offer(next[1] - next[0])
-                    sum = sum - (if (queue.size > k) queue.poll() else 0)
-                }
+        val pq = PriorityQueue<Long?>()
+        for (e in adj[v]) {
+            val w = if (e[0] == v) e[1] else e[0]
+            if (w == parent) {
+                continue
             }
+            val res = dfs(w, v)
+            val max = max((e[2] + res[0]), res[1])
+            sum += max
+            pq.add(max - res[1])
+        }
+        val res = LongArray(2)
+        while (pq.size > k) {
+            sum -= pq.poll()!!
+        }
+        res[1] = sum
+        while (pq.size > k - 1) {
+            sum -= pq.poll()!!
         }
-        return longArrayOf(sum, sum - (if (queue.size < k) 0 else queue.peek())!!)
+        res[0] = sum
+        return res
     }
 }