2302. Count Subarrays With Score Less Than K Hard

1/**
2 * [2302] Count Subarrays With Score Less Than K
3 *
4 * The score of an array is defined as the product of its sum and its length.
5 *
6 * For example, the score of [1, 2, 3, 4, 5] is (1 + 2 + 3 + 4 + 5) * 5 = 75.
7 *
8 * 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.
9 * A subarray is a contiguous sequence of elements within an array.
10 *  
11 * Example 1:
12 *
13 * Input: nums = [2,1,4,3,5], k = 10
14 * Output: 6
15 * Explanation:
16 * The 6 subarrays having scores less than 10 are:
17 * - [2] with score 2 * 1 = 2.
18 * - [1] with score 1 * 1 = 1.
19 * - [4] with score 4 * 1 = 4.
20 * - [3] with score 3 * 1 = 3.
21 * - [5] with score 5 * 1 = 5.
22 * - [2,1] with score (2 + 1) * 2 = 6.
23 * 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.
24 * Example 2:
25 *
26 * Input: nums = [1,1,1], k = 5
27 * Output: 5
28 * Explanation:
29 * Every subarray except [1,1,1] has a score less than 5.
30 * [1,1,1] has a score (1 + 1 + 1) * 3 = 9, which is greater than 5.
31 * Thus, there are 5 subarrays having scores less than 5.
32 *
33 *  
34 * Constraints:
35 *
36 * 1 <= nums.length <= 10^5
37 * 1 <= nums[i] <= 10^5
38 * 1 <= k <= 10^15
39 *
40 */
41pub struct Solution {}
42
43// problem: https://leetcode.com/problems/count-subarrays-with-score-less-than-k/
44// discuss: https://leetcode.com/problems/count-subarrays-with-score-less-than-k/discuss/?currentPage=1&orderBy=most_votes&query=
45
46// submission codes start here
47
48impl Solution {
49    pub fn count_subarrays(nums: Vec<i32>, k: i64) -> i64 {
50        let mut result = 0i64;
51        let mut sum = 0i64;
52        let mut start = 0;
53        for end in 0..nums.len() {
54            sum += nums[end] as i64;
55            while sum as i64 * (end - start + 1) as i64 >= k {
56                sum -= nums[start] as i64;
57                start += 1;
58            }
59            result += (end - start + 1) as i64;
60        }
61
62        result as i64
63    }
64}
65
66// submission codes end
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn test_2302() {
74        assert_eq!(Solution::count_subarrays(vec![2, 1, 4, 3, 5], 10), 6);
75    }
76}
77