Leetcode: 2444. Count Subarrays With Fixed Bounds

Problem Statement

Understand the problem

Given an array \(a\) and numbers \(minK\) and \(maxK\), return the number of subarrays \(s\) where \(min(s)=minK\) and \(max(s)=maxK\).

Devise a plan

If \(a[i] < minK\) or \(a[i] > maxK\), then no subarray that contains \(a[i]\) should be counted. So, we can split the given array in subarray that contains only elements between \(minK\) and \(maxK\). For now on, we suppose that \(a\) only contains elements from \(minK\) to \(maxK\). Pattern: Count subarrays starting/ending in a position. For each position \(i\), we find the smallest index \(j\) where \(a[i..j]\) contains at least one value \(minK\) and one value \(maxK\). The number of subarrays starting on \(i\) is going to be \(|a|-j\). Time complexity and space \(O(n)\).

Carry out the plan

class Solution:
    def countSubarrays(self, nums: List[int], minK: int, maxK: int) -> int:
        ans = 0
        for t, g in groupby(nums, key=lambda n: minK <= n <= maxK):
            if not t:
                continue
            group = list(g)
            lastMin = -1
            lastMax = -1
            for i in range(len(group) - 1, -1, -1):
                if group[i] == minK:
                    lastMin = i
                if group[i] == maxK:
                    lastMax = i
                if lastMin != -1 and lastMax != -1:
                    j = max(lastMin, lastMax)
                    ans += len(group) - j
        return ans