3478.choose-k-elements-with-maximum-sum.md

3478. Choose K Elements With Maximum Sum

TODO: Understand and AC this problem.

Given a index i, how to find all j? For i, find all values of j < nums[i].

pair with (value, index):

value 4, 2, 1, 5, 3
index 0  1  2  3  4

sort by value:
value 1, 2, 3, 4, 5
index 2  1  4  0  3
      i
      |  i
      |--|  i
      |-----|  i
      |--------|  i
      take most k to maximum sum

Brute Force (TLE)

fun findMaxSum(nums1: IntArray, nums2: IntArray, k: Int): LongArray {
    val n = nums1.size
    val valueIndex = Array(n) {
        Pair<Int, Int>(0, 0)
    }
    for (i in nums1.indices) {
        valueIndex[i] = nums1[i] to i
    }
    valueIndex.sortBy { it.first }
    val ans = LongArray(n)
    for (i in 1 until valueIndex.size) {
        val (value, index) = valueIndex[i]
        val minHeap = PriorityQueue<Int>()
        for (j in 0 until i) {
            val (value2, index2) = valueIndex[j]
            if (value2 < value) {
                minHeap.add(nums2[index2])
                if (minHeap.size > k) {
                    minHeap.poll()
                }
            }
        }
        var sum = 0L
        while (minHeap.isNotEmpty()) sum += minHeap.poll().toLong()
        ans[index] = sum
    }
    return ans
}

Heap + Prefix Sum

fun findMaxSum(nums1: IntArray, nums2: IntArray, k: Int): LongArray {
    val n = nums1.size
    val valueIndex = Array(n) {
        Pair<Int, Int>(0, 0)
    }
    for (i in nums1.indices) {
        valueIndex[i] = nums1[i] to i
    }
    valueIndex.sortBy { it.first }
    val minHeap = PriorityQueue<Int>()
    val map = HashMap<Int, Long>()
    var sum = 0L
    var j = 0
    val ans = LongArray(n)
    for (i in valueIndex.indices) {
        val (value, index) = valueIndex[i]
        while (j < i) {
            val (value2, index2) = valueIndex[j]
            if (value <= value2) break

            minHeap.add(nums2[index2])
            sum += nums2[index2].toLong()
            if (minHeap.size > k) {
                sum -= minHeap.poll().toLong()
            }
            j++
        }
        map[index] = sum
    }
    for (i in valueIndex.indices) {
        ans[i] = map[i] ?: 0L
    }
    return ans
}

Analysis

Key Differences Between Solutions

1. Time Complexity:

   • Brute Force: O(n² * k * log k) - For each element, we process all previous elements and maintain a heap
   • Heap + Prefix Sum: O(n * log k) - We process each element only once and maintain a single heap

2. Space Complexity:

   • Brute Force: O(k) for each iteration, creating new heaps
   • Heap + Prefix Sum: O(n) for storing the map and O(k) for the single heap

3. Processing Approach:

   • Brute Force: For each element, we create a new heap and process all previous elements
   • Heap + Prefix Sum: We maintain a single heap and process elements in order, reusing the sum

Why the Second Solution Works Better

1. Single Pass Processing:

   • Instead of creating a new heap for each element, we maintain a single heap
   • We process elements in sorted order, which allows us to reuse the sum from previous calculations

2. Efficient Sum Maintenance:

   • We keep track of the running sum using a single variable
   • When the heap size exceeds k, we remove the smallest element and subtract it from the sum
   • This avoids recalculating the sum for each element

3. Smart Index Tracking:

   • We use a map to store the sum for each index
   • This allows us to process elements in sorted order but maintain the original order in the result

4. Early Break Condition:

   • The solution uses if (value <= value2) break to stop processing when we encounter a larger value
   • This optimization prevents unnecessary processing of elements that won't contribute to the sum

The second solution is more efficient because it eliminates redundant calculations and reuses the work done for previous elements, making it suitable for larger input sizes.

Key Intuitions and Observations

1. Redundant Calculations:

   • In the brute force solution, we were recalculating the same sums multiple times
   • For each element i, we were processing all previous elements (0 to i-1) again
   • This led to O(n²) complexity as we were doing the same work repeatedly

2. Order Matters:

   • We noticed that if we process elements in sorted order of nums1, we can maintain a single heap
   • When we process a new element, we only need to consider elements that are smaller than the current value
   • This is why we sort the valueIndex array by nums1 values

3. Running Sum Pattern:

   • Instead of recalculating the sum for each element, we can maintain a running sum
   • When we add a new element to the heap, we add its value to the sum
   • When we remove an element (if heap size > k), we subtract its value from the sum
   • This gives us O(1) sum updates instead of O(k) recalculations

4. State Reuse:

   • The key insight was that we don't need to create a new heap for each element
   • The state (heap and sum) from processing previous elements can be reused
   • We just need to maintain the k largest elements from nums2 that correspond to smaller nums1 values

5. Early Termination:

   • We observed that once we encounter a nums1 value that's greater than or equal to the current value, we can stop processing
   • This is because those elements won't contribute to the sum for the current element
   • This optimization is implemented in the if (value <= value2) break condition

These observations led to the transformation from a brute force approach to an efficient solution that processes each element only once while maintaining the correct order of results.