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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
| // @lc code=start
impl Solution {
/// ## 解题思路
/// - 前缀和+归并排序
pub fn count_range_sum(nums: Vec<i32>, lower: i32, upper: i32) -> i32 {
// 计算前缀和
let mut prefix_sums = vec![0_i64];
for &n in &nums {
let last = *prefix_sums.last().unwrap_or(&0);
prefix_sums.push(n as i64 + last);
}
fn merge(
prefix_sums: &mut Vec<i64>,
left: usize,
right: usize,
res: &mut i64,
lower: i64,
upper: i64,
) {
if right - left <= 1 {
return;
}
let mid = (left + right) >> 1;
merge(prefix_sums, left, mid, res, lower, upper);
merge(prefix_sums, mid, right, res, lower, upper);
// 计算区间和个数
let (mut s, mut e) = (mid, mid);
for l in left..mid {
while s < right && prefix_sums[s] - prefix_sums[l] < lower {
s += 1;
}
while e < right && prefix_sums[e] - prefix_sums[l] <= upper {
e += 1;
}
*res += (e - s) as i64;
}
// 合并
let left_sums = prefix_sums[left..mid].to_vec();
let right_sums = prefix_sums[mid..right].to_vec();
let (mut l, mut r) = (0, 0);
let mut i = left;
while l < left_sums.len() && r < right_sums.len() {
if left_sums[l] < right_sums[r] {
prefix_sums[i] = left_sums[l];
l += 1;
} else {
prefix_sums[i] = right_sums[r];
r += 1;
}
i += 1;
}
while l < left_sums.len() {
prefix_sums[i] = left_sums[l];
l += 1;
i += 1;
}
while r < right_sums.len() {
prefix_sums[i] = right_sums[r];
r += 1;
i += 1;
}
}
let n = prefix_sums.len();
let mut res = 0;
merge(&mut prefix_sums, 0, n, &mut res, lower as i64, upper as i64);
res as i32
}
}
// @lc code=end
|