Skip to content

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

Thinking

A subarray’s score is its sum times its length. Enumerating all subarrays is \(O(n^2)\) and fails for \(n \le 10^5\). All values are positive, so for a fixed right end a longer window has a larger score, and valid left ends form a prefix.

Prefix sums answer any range sum in constant time. Binary-search the largest length \(l\) with \((s[i]-s[i-l])\times l < k\). That right end contributes \(l\) subarrays.

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}\).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
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;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
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
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
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;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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

Thinking

Method 1 stores a prefix array and pays a logarithmic cost per right end. Positivity makes the score decrease as the left end moves right, so the left pointer only advances. Maintain the window sum and count valid subarrays in one linear scan.

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)\).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
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;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
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;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
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
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
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;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
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
    }
}

Comments