LeetCode in Kotlin

3318. Find X-Sum of All K-Long Subarrays I

Easy

You are given an array nums of n integers and two integers k and x.

The x-sum of an array is calculated by the following procedure:

Note that if an array has less than x distinct elements, its x-sum is the sum of the array.

Return an integer array answer of length n - k + 1 where answer[i] is the x-sum of the subarray nums[i..i + k - 1].

Example 1:

Input: nums = [1,1,2,2,3,4,2,3], k = 6, x = 2

Output: [6,10,12]

Explanation:

Example 2:

Input: nums = [3,8,7,8,7,5], k = 2, x = 2

Output: [11,15,15,15,12]

Explanation:

Since k == x, answer[i] is equal to the sum of the subarray nums[i..i + k - 1].

Constraints:

Solution

import java.util.Comparator
import java.util.HashMap
import java.util.PriorityQueue

class Solution {
    private class Pair(num: Int, freq: Int) {
        var num: Int
        var freq: Int

        init {
            this.num = num
            this.freq = freq
        }
    }

    fun findXSum(nums: IntArray, k: Int, x: Int): IntArray {
        val n = nums.size
        val ans = IntArray(n - k + 1)
        for (i in 0 until n - k + 1) {
            val map = HashMap<Int, Int>()
            val pq =
                PriorityQueue<Pair>(
                    Comparator { a: Pair, b: Pair ->
                        if (a.freq == b.freq) {
                            return@Comparator b.num - a.num
                        }
                        b.freq - a.freq
                    },
                )
            for (j in i until i + k) {
                map.put(nums[j], map.getOrDefault(nums[j], 0) + 1)
            }
            for (entry in map.entries) {
                pq.add(Pair(entry.key, entry.value))
            }
            var count = x
            var sum = 0
            while (pq.isNotEmpty() && count > 0) {
                val pair = pq.remove()
                sum += pair.num * pair.freq
                count--
            }
            ans[i] = sum
        }
        return ans
    }
}