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 <= 1051 <= nums[i] <= 1051 <= k <= 1015
Solutions
Solution 1: Prefix Sum + Binary Search
First, we calculate the prefix sum array $s$ of the array $\textit{nums}$, where $s[i]$ represents the sum of the first $i$ elements of $\textit{nums}$.
Next, we enumerate each element of $\textit{nums}$ as the last element of a subarray. For each element, we can use binary search to find the maximum length $l$ such that $s[i] - s[i - l] \times l < k$. The number of subarrays ending at this element is $l$, and summing up all $l$ gives the final answer.
The time complexity is $O(n \times \log n)$, and the space complexity is $O(n)$, where $n$ is the length of the array $\textit{nums}$.
Python3
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)):
l, r = 0, i
while l < r:
mid = (l + r + 1) >> 1
if (s[i] - s[i - mid]) * mid < k:
l = mid
else:
r = mid - 1
ans += l
return ans
Java
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 l = 0, r = i;
while (l < r) {
int mid = (l + r + 1) >> 1;
if ((s[i] - s[i - mid]) * mid < k) {
l = mid;
} else {
r = mid - 1;
}
}
ans += l;
}
return ans;
}
}
C++
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 l = 0, r = i;
while (l < r) {
int mid = (l + r + 1) >> 1;
if ((s[i] - s[i - mid]) * mid < k) {
l = mid;
} else {
r = mid - 1;
}
}
ans += l;
}
return ans;
}
};
Go
func countSubarrays(nums []int, k int64) (ans int64) {
n := len(nums)
s := make([]int64, n+1)
for i, x := range nums {
s[i+1] = s[i] + int64(x)
}
for i := 1; i <= n; i++ {
l, r := 0, i
for l < r {
mid := (l + r + 1) >> 1
if (s[i]-s[i-mid])*int64(mid) < k {
l = mid
} else {
r = mid - 1
}
}
ans += int64(l)
}
return
}
TypeScript
function countSubarrays(nums: number[], k: number): number {
const n = nums.length;
const s: number[] = Array(n + 1).fill(0);
for (let i = 0; i < n; ++i) {
s[i + 1] = s[i] + nums[i];
}
let ans = 0;
for (let i = 1; i <= n; ++i) {
let [l, r] = [0, i];
while (l < r) {
const mid = (l + r + 1) >> 1;
if ((s[i] - s[i - mid]) * mid < k) {
l = mid;
} else {
r = mid - 1;
}
}
ans += l;
}
return ans;
}
Rust
impl Solution {
pub fn count_subarrays(nums: Vec<i32>, k: i64) -> i64 {
let n = nums.len();
let mut s = vec![0i64; n + 1];
for i in 0..n {
s[i + 1] = s[i] + nums[i] as i64;
}
let mut ans = 0i64;
for i in 1..=n {
let mut l = 0;
let mut r = i;
while l < r {
let mid = (l + r + 1) / 2;
let sum = s[i] - s[i - mid];
if sum * (mid as i64) < k {
l = mid;
} else {
r = mid - 1;
}
}
ans += l as i64;
}
ans
}
}
Solution 2: Two Pointers
We can use the two-pointer technique to maintain a sliding window such that the sum of elements in the window is less than $k$. The number of subarrays ending at the current element is equal to the length of the window. Summing up all the window lengths gives the final answer.
The time complexity is $O(n)$, where $n$ is the length of the array $\textit{nums}$. The space complexity is $O(1)$.
Python3
class Solution:
def countSubarrays(self, nums: List[int], k: int) -> int:
ans = s = j = 0
for i, x in enumerate(nums):
s += x
while s * (i - j + 1) >= k:
s -= nums[j]
j += 1
ans += i - j + 1
return ans
Java
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;
}
}
C++
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;
}
};
Go
func countSubarrays(nums []int, k int64) (ans int64) {
s, j := 0, 0
for i, x := range nums {
s += x
for int64(s*(i-j+1)) >= k {
s -= nums[j]
j++
}
ans += int64(i - j + 1)
}
return
}
TypeScript
function countSubarrays(nums: number[], k: number): number {
let [ans, s, j] = [0, 0, 0];
for (let i = 0; i < nums.length; ++i) {
s += nums[i];
while (s * (i - j + 1) >= k) {
s -= nums[j++];
}
ans += i - j + 1;
}
return ans;
}
Rust
impl Solution {
pub fn count_subarrays(nums: Vec<i32>, k: i64) -> i64 {
let mut ans = 0i64;
let mut s = 0i64;
let mut j = 0;
for i in 0..nums.len() {
s += nums[i] as i64;
while s * (i as i64 - j as i64 + 1) >= k {
s -= nums[j] as i64;
j += 1;
}
ans += i as i64 - j as i64 + 1;
}
ans
}
}