Leetcode: 2302. Count Subarrays With Score Less Than K

Problem Statement: Given an array of integers \(a\) and an integer \(k\), return the number of non-empty subarrays \(a[i..j]\) such that \(s(i,j)=(a[i]+a[i+1]+...+a[j]) \times (j - i + 1) < k\).

Be \(i\) and \(j\) such \(s(i,j)How can we extend the solution for \(i\) to \(i+1\)? For \(i+1\), we want to find the greatest index \(j'\) such that \(s(i+1, j')One-dimensional Prefix Sum.

Time complexity: \(O(|a|)\).
Space complexity: \(O(1)\).

from collections import deque
from typing import List


def solve(nums, k):
    n = len(nums)

    pref = [0] * n
    pref[0] = nums[0]
    for i in range(1, n):
        pref[i] = nums[i] + pref[i - 1]

    def score(i, j):
        size = j - i + 1
        return (pref[j] - (0 if i == 0 else pref[i - 1])) * size

    j = n - 1
    ans = 0
    for i in range(n - 1, -1, -1):
        while i <= j and score(i, j) >= k:
            j -= 1
        if i <= j:
            ans += j - i + 1

    return ans


assert solve([2, 1, 4, 3, 5], 10) == 6
assert solve([1, 1, 1], 5) == 5
assert solve([3, 3], 2) == 0
assert solve([1, 1, 1, 1], 5) == 7
assert solve([5, 2, 6, 8, 9, 7], 50) == 13


class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        return solve(nums, k)