Count Subarrays Where Max Element Appears at Least K Times - LeetCode Daily Challenge

Problem Statement

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.

Try here before watching the video.

Video Solution

Java Code

class Solution {
    public long countSubarrays(int[] nums, int k) {
        int start = 0, maxElement = 0, count = 0;

        long subarrayCount = 0;

        for(int num : nums) {
            maxElement = Math.max(maxElement, num);
        }

        for(int end = 0; end < nums.length; end++) {
            int curr = nums[end];
            if(curr == maxElement) {
                count += 1;
            }

            while(count == k) {
                if(nums[start] == maxElement) {
                    count -= 1;
                }
                start += 1;
            }

            subarrayCount += start;
        }

        return subarrayCount;
    }
}

C++ Code

#include <vector>
using namespace std;

class Solution {
public:
    long long countSubarrays(vector<int>& nums, int k) {
        int start = 0, maxElement = 0, count = 0;
        long long subarrayCount = 0;

        for (int num : nums) {
            maxElement = max(maxElement, num);
        }

        for (int end = 0; end < nums.size(); end++) {
            int curr = nums[end];
            if (curr == maxElement) {
                count += 1;
            }

            while (count == k) {
                if (nums[start] == maxElement) {
                    count -= 1;
                }
                start += 1;
            }

            subarrayCount += start;
        }

        return subarrayCount;
    }
};

Python Code

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        start = 0
        max_element = max(nums)
        count = 0
        subarray_count = 0

        for end in range(len(nums)):
            curr = nums[end]
            if curr == max_element:
                count += 1

            while count == k:
                if nums[start] == max_element:
                    count -= 1
                start += 1

            subarray_count += start

        return subarray_count

Javascript Code

var countSubarrays = function(nums, k) {
    let start = 0;
    let maxElement = Math.max(...nums);
    let count = 0;
    let subarrayCount = 0;

    for (let end = 0; end < nums.length; end++) {
        const curr = nums[end];
        if (curr === maxElement) {
            count++;
        }

        while (count === k) {
            if (nums[start] === maxElement) {
                count--;
            }
            start++;
        }

        subarrayCount += start;
    }

    return subarrayCount;
};

Go Code

func countSubarrays(nums []int, k int) int64 {
    start := 0
    maxElement := nums[0]
    for _, num := range nums {
        if num > maxElement {
            maxElement = num
        }
    }

    count := 0
    var subarrayCount int64 = 0

    for end := 0; end < len(nums); end++ {
        curr := nums[end]
        if curr == maxElement {
            count++
        }

        for count == k {
            if nums[start] == maxElement {
                count--
            }
            start++
        }

        subarrayCount += int64(start)
    }

    return subarrayCount
}

Complexity Analysis

Time Complexity: O(N)
Space Complexity: O(1)