LeetCode in Kotlin

3974. Maximum Total Sum of K Selected Elements

Medium

You are given an integer array nums and two integers k and mul.

Select exactly k elements from nums. Process these elements one by one in any order you choose.

For each selected element, independently choose one of the following:

• Add the element’s value to the total sum, or
• Multiply the element by the current value of mul and add the result to the total sum.

After processing each selected element, mul decreases by 1, regardless of which option was chosen. The current value of mul may become 0 or negative.

Return an integer denoting the maximum possible total sum.

Example 1:

Input: nums = [6,1,2,9], k = 3, mul = 2

Output: 26

Explanation:

One optimal way:

• One optimal selection is nums[3] = 9, nums[0] = 6, and nums[2] = 2.
• Process nums[3] = 9 first: choose multiplication, so it contributes 9 * 2 = 18. Now, mul becomes 1.
• Process nums[0] = 6 next: choose multiplication, so it contributes 6 * 1 = 6. Now, mul becomes 0.
• Process nums[2] = 2 last: choose addition, so it contributes 2.
• The total sum is 18 + 6 + 2 = 26.

Example 2:

Input: nums = [3,7,5,2], k = 2, mul = 4

Output: 43

Explanation:

One optimal way:

• One optimal selection is nums[1] = 7 and nums[2] = 5.
• Process nums[1] = 7 first: choose multiplication, so it contributes 7 * 4 = 28. Now, mul becomes 3.
• Process nums[2] = 5 next: choose multiplication, so it contributes 5 * 3 = 15.
• The total sum is 28 + 15 = 43.

Example 3:

Input: nums = [4,4], k = 1, mul = 1

Output: 4

Explanation:

One optimal way:

• One optimal selection is nums[0] = 4.
• Process nums[0] = 4: choose multiplication, so it contributes 4 * 1 = 4.
• The total sum is 4.

Constraints:

• 1 <= nums.length <= 105
• 1 <= nums[i] <= 105
• 1 <= k <= nums.length
• 1 <= mul <= 105

Solution

import kotlin.math.max
import kotlin.math.min

class Solution {
    fun maxSum(nums: IntArray, k: Int, mul: Int): Long {
        var k = k
        var mul = mul
        val count = IntArray(MAX)
        var max = nums[0]
        var sum: Long = 0
        for (n in nums) {
            max = max(max, n)
            count[n]++
        }
        while (max >= 0 && k > 0) {
            if (count[max] != 0) {
                var min = min(mul, count[max])
                min = min(k, min)
                k -= min
                mul -= min
                val local =
                    ((mul + min).toLong() * (mul + min + 1)) / 2 - (mul.toLong() * (mul + 1)) / 2
                sum += local * max
                if (mul == 0) {
                    count[max] -= min
                    break
                }
            }
            max--
        }
        while (max >= 0 && k > 0) {
            if (count[max] != 0) {
                val min = min(k, count[max])
                k -= min
                sum += max.toLong() * min
            }
            max--
        }
        return sum
    }

    companion object {
        private val MAX = 100000 + 1
    }
}

This site is open source. Improve this page.