Question

2302. Count Subarrays With Score Less Than K

中文文档

Description

The score of an array is defined as the product of its sum and its length.

• For example, the score of [1, 2, 3, 4, 5] is (1 + 2 + 3 + 4 + 5) * 5 = 75.

Given a positive integer array nums and an integer k, return _the number of non-empty subarrays of_ nums _whose score is strictly less than_ k.

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

 

Example 1:

Input: nums = [2,1,4,3,5], k = 10
Output: 6
Explanation:
The 6 subarrays having scores less than 10 are:
- [2] with score 2 * 1 = 2.
- [1] with score 1 * 1 = 1.
- [4] with score 4 * 1 = 4.
- [3] with score 3 * 1 = 3. 
- [5] with score 5 * 1 = 5.
- [2,1] with score (2 + 1) * 2 = 6.
Note that subarrays such as [1,4] and [4,3,5] are not considered because their scores are 10 and 36 respectively, while we need scores strictly less than 10.

Example 2:

Input: nums = [1,1,1], k = 5
Output: 5
Explanation:
Every subarray except [1,1,1] has a score less than 5.
[1,1,1] has a score (1 + 1 + 1) * 3 = 9, which is greater than 5.
Thus, there are 5 subarrays having scores less than 5.

 

Constraints:

• 1 <= nums.length <= 105
• 1 <= nums[i] <= 105
• 1 <= k <= 1015

Solutions

Solution 1: Prefix Sum + Binary Search

First, we calculate the prefix sum array $s$ of the array $nums$, where $s[i]$ represents the sum of the first $i$ elements of the array $nums$.

Next, we enumerate each element of the array $nums$ as the last element of the subarray. For each element, we can find the maximum length $l$ such that $s[i] - s[i - l] \times l < k$ by binary search. The number of subarrays with this element as the last element is $l$, and we add all $l$ to get the answer.

The time complexity is $O(n \times \log n)$, and the space complexity is $O(n)$. Here, $n$ is the length of the array $nums$.

class Solution:
  def countSubarrays(self, nums: List[int], k: int) -> int:
    s = list(accumulate(nums, initial=0))
    ans = 0
    for i in range(1, len(s)):
      left, right = 0, i
      while left < right:
        mid = (left + right + 1) >> 1
        if (s[i] - s[i - mid]) * mid < k:
          left = mid
        else:
          right = mid - 1
      ans += left
    return ans

class Solution {
  public long countSubarrays(int[] nums, long k) {
    int n = nums.length;
    long[] s = new long[n + 1];
    for (int i = 0; i < n; ++i) {
      s[i + 1] = s[i] + nums[i];
    }
    long ans = 0;
    for (int i = 1; i <= n; ++i) {
      int left = 0, right = i;
      while (left < right) {
        int mid = (left + right + 1) >> 1;
        if ((s[i] - s[i - mid]) * mid < k) {
          left = mid;
        } else {
          right = mid - 1;
        }
      }
      ans += left;
    }
    return ans;
  }
}

class Solution {
public:
  long long countSubarrays(vector<int>& nums, long long k) {
    int n = nums.size();
    long long s[n + 1];
    s[0] = 0;
    for (int i = 0; i < n; ++i) {
      s[i + 1] = s[i] + nums[i];
    }
    long long ans = 0;
    for (int i = 1; i <= n; ++i) {
      int left = 0, right = i;
      while (left < right) {
        int mid = (left + right + 1) >> 1;
        if ((s[i] - s[i - mid]) * mid < k) {
          left = mid;
        } else {
          right = mid - 1;
        }
      }
      ans += left;
    }
    return ans;
  }
};

func countSubarrays(nums []int, k int64) (ans int64) {
  n := len(nums)
  s := make([]int64, n+1)
  for i, v := range nums {
    s[i+1] = s[i] + int64(v)
  }
  for i := 1; i <= n; i++ {
    left, right := 0, i
    for left < right {
      mid := (left + right + 1) >> 1
      if (s[i]-s[i-mid])*int64(mid) < k {
        left = mid
      } else {
        right = mid - 1
      }
    }
    ans += int64(left)
  }
  return
}

Solution 2: Two Pointers

We can use two pointers to maintain a sliding window, so that the sum of the elements in the window is less than $k$. The number of subarrays with the current element as the last element is the length of the window, and we add all window lengths to get the answer.

The time complexity is $O(n)$, where $n$ is the length of the array $nums$. The space complexity is $O(1)$.

class Solution:
  def countSubarrays(self, nums: List[int], k: int) -> int:
    ans = s = j = 0
    for i, v in enumerate(nums):
      s += v
      while s * (i - j + 1) >= k:
        s -= nums[j]
        j += 1
      ans += i - j + 1
    return ans

class Solution {
  public long countSubarrays(int[] nums, long k) {
    long ans = 0, s = 0;
    for (int i = 0, j = 0; i < nums.length; ++i) {
      s += nums[i];
      while (s * (i - j + 1) >= k) {
        s -= nums[j++];
      }
      ans += i - j + 1;
    }
    return ans;
  }
}

class Solution {
public:
  long long countSubarrays(vector<int>& nums, long long k) {
    long long ans = 0, s = 0;
    for (int i = 0, j = 0; i < nums.size(); ++i) {
      s += nums[i];
      while (s * (i - j + 1) >= k) {
        s -= nums[j++];
      }
      ans += i - j + 1;
    }
    return ans;
  }
};

func countSubarrays(nums []int, k int64) (ans int64) {
  s, j := 0, 0
  for i, v := range nums {
    s += v
    for int64(s*(i-j+1)) >= k {
      s -= nums[j]
      j++
    }
    ans += int64(i - j + 1)
  }
  return
}