|
| 1 | +[](https://github.com/javadev/LeetCode-in-Kotlin) |
| 2 | +[](https://github.com/javadev/LeetCode-in-Kotlin/fork) |
| 3 | + |
| 4 | +## 3478\. Choose K Elements With Maximum Sum |
| 5 | + |
| 6 | +Medium |
| 7 | + |
| 8 | +You are given two integer arrays, `nums1` and `nums2`, both of length `n`, along with a positive integer `k`. |
| 9 | + |
| 10 | +For each index `i` from `0` to `n - 1`, perform the following: |
| 11 | + |
| 12 | +* Find **all** indices `j` where `nums1[j]` is less than `nums1[i]`. |
| 13 | +* Choose **at most** `k` values of `nums2[j]` at these indices to **maximize** the total sum. |
| 14 | + |
| 15 | +Return an array `answer` of size `n`, where `answer[i]` represents the result for the corresponding index `i`. |
| 16 | + |
| 17 | +**Example 1:** |
| 18 | + |
| 19 | +**Input:** nums1 = [4,2,1,5,3], nums2 = [10,20,30,40,50], k = 2 |
| 20 | + |
| 21 | +**Output:** [80,30,0,80,50] |
| 22 | + |
| 23 | +**Explanation:** |
| 24 | + |
| 25 | +* For `i = 0`: Select the 2 largest values from `nums2` at indices `[1, 2, 4]` where `nums1[j] < nums1[0]`, resulting in `50 + 30 = 80`. |
| 26 | +* For `i = 1`: Select the 2 largest values from `nums2` at index `[2]` where `nums1[j] < nums1[1]`, resulting in 30. |
| 27 | +* For `i = 2`: No indices satisfy `nums1[j] < nums1[2]`, resulting in 0. |
| 28 | +* For `i = 3`: Select the 2 largest values from `nums2` at indices `[0, 1, 2, 4]` where `nums1[j] < nums1[3]`, resulting in `50 + 30 = 80`. |
| 29 | +* For `i = 4`: Select the 2 largest values from `nums2` at indices `[1, 2]` where `nums1[j] < nums1[4]`, resulting in `30 + 20 = 50`. |
| 30 | + |
| 31 | +**Example 2:** |
| 32 | + |
| 33 | +**Input:** nums1 = [2,2,2,2], nums2 = [3,1,2,3], k = 1 |
| 34 | + |
| 35 | +**Output:** [0,0,0,0] |
| 36 | + |
| 37 | +**Explanation:** |
| 38 | + |
| 39 | +Since all elements in `nums1` are equal, no indices satisfy the condition `nums1[j] < nums1[i]` for any `i`, resulting in 0 for all positions. |
| 40 | + |
| 41 | +**Constraints:** |
| 42 | + |
| 43 | +* `n == nums1.length == nums2.length` |
| 44 | +* <code>1 <= n <= 10<sup>5</sup></code> |
| 45 | +* <code>1 <= nums1[i], nums2[i] <= 10<sup>6</sup></code> |
| 46 | +* `1 <= k <= n` |
| 47 | + |
| 48 | +## Solution |
| 49 | + |
| 50 | +```kotlin |
| 51 | +import java.util.PriorityQueue |
| 52 | + |
| 53 | +class Solution { |
| 54 | + fun findMaxSum(nums1: IntArray, nums2: IntArray, k: Int): LongArray { |
| 55 | + val n = nums1.size |
| 56 | + val ans = LongArray(n) |
| 57 | + val ps = Array<Point>(n) { i -> Point(nums1[i], nums2[i], i) } |
| 58 | + ps.sortWith { p1: Point, p2: Point -> p1.x.compareTo(p2.x) } |
| 59 | + val pq = PriorityQueue<Int?>() |
| 60 | + var s: Long = 0 |
| 61 | + var i = 0 |
| 62 | + while (i < n) { |
| 63 | + var j = i |
| 64 | + while (j < n && ps[j].x == ps[i].x) { |
| 65 | + ans[ps[j].i] = s |
| 66 | + j++ |
| 67 | + } |
| 68 | + for (p in i..<j) { |
| 69 | + val cur = ps[p].y |
| 70 | + if (pq.size < k) { |
| 71 | + pq.offer(cur) |
| 72 | + s += cur.toLong() |
| 73 | + } else if (cur > pq.peek()!!) { |
| 74 | + s -= pq.poll()!!.toLong() |
| 75 | + pq.offer(cur) |
| 76 | + s += cur.toLong() |
| 77 | + } |
| 78 | + } |
| 79 | + i = j |
| 80 | + } |
| 81 | + return ans |
| 82 | + } |
| 83 | + |
| 84 | + private class Point(var x: Int, var y: Int, var i: Int) |
| 85 | +} |
| 86 | +``` |
0 commit comments