LeetCode in Kotlin

2962. Count Subarrays Where Max Element Appears at Least K Times

Medium

You are given an integer array nums and a positive integer k.

Return the number of subarrays where the maximum element of nums appears at least k times in that subarray.

A subarray is a contiguous sequence of elements within an array.

Example 1:

Input: nums = [1,3,2,3,3], k = 2

Output: 6

Explanation: The subarrays that contain the element 3 at least 2 times are: [1,3,2,3], [1,3,2,3,3], [3,2,3], [3,2,3,3], [2,3,3] and [3,3].

Example 2:

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

Output: 0

Explanation: No subarray contains the element 4 at least 3 times.

Constraints:

Solution

class Solution {
    fun countSubarrays(nums: IntArray, k: Int): Long {
        val st = IntArray(nums.size + 1)
        var si = 0
        var m = 0
        for (i in nums.indices) {
            if (m < nums[i]) {
                m = nums[i]
                si = 0
            }
            if (m == nums[i]) {
                st[si++] = i
            }
        }
        if (si < k) {
            return 0
        }
        var r: Long = 0
        st[si] = nums.size
        for (i in k..si) {
            r += (st[i - k] + 1).toLong() * (st[i] - st[i - 1])
        }
        return r
    }
}