1#include <vector>
2#include <algorithm>
3#include <map>
4
5// Fenwick Tree (BIT) implementation
6class FenwickTree {
7private:
8 std::vector<int> bit;
9 int size;
10
11public:
12 FenwickTree(int s) : size(s), bit(s + 1, 0) {}
13
14 void update(int idx, int delta) {
15 idx++; // 1-based indexing
16 while (idx <= size) {
17 bit[idx] += delta;
18 idx += idx & (-idx);
19 }
20 }
21
22 int query(int idx) {
23 idx++; // 1-based indexing
24 int sum = 0;
25 while (idx > 0) {
26 sum += bit[idx];
27 idx -= idx & (-idx);
28 }
29 return sum;
30 }
31
32 // Query sum in range [l, r]
33 int query_range(int l, int r) {
34 if (l > r) return 0;
35 return query(r) - query(l - 1);
36 }
37};
38
39long long countSubarrays(const std::vector<int>& nums, int L, int R) {
40 int n = nums.size();
41 std::vector<long long> prefix_sums(n + 1, 0);
42 for (int i = 0; i < n; ++i) {
43 prefix_sums[i + 1] = prefix_sums[i] + nums[i];
44 }
45
46 // Collect all values relevant for coordinate compression
47 // This includes prefix sums themselves and potential query bounds (P[j]-L, P[j]-R)
48 std::vector<long long> values_to_compress;
49 for (long long p_sum : prefix_sums) {
50 values_to_compress.push_back(p_sum);
51 values_to_compress.push_back(p_sum - L);
52 values_to_compress.push_back(p_sum - R);
53 }
54 std::sort(values_to_compress.begin(), values_to_compress.end());
55 values_to_compress.erase(std::unique(values_to_compress.begin(), values_to_compress.end()), values_to_compress.end());
56
57 std::map<long long, int> compressed_coords;
58 for (int i = 0; i < values_to_compress.size(); ++i) {
59 compressed_coords[values_to_compress[i]] = i;
60 }
61
62 FenwickTree ft(values_to_compress.size());
63 long long count = 0;
64
65 // The initial prefix sum P[0] (which is 0) contributes to subsequent calculations
66 // For P[0], it implies an empty prefix before index 0.
67 ft.update(compressed_coords[0], 1);
68
69 for (int j = 1; j <= n; ++j) {
70 long long current_p_sum = prefix_sums[j];
71
72 // Find the range [current_p_sum - R, current_p_sum - L]
73 // Get compressed indices for the lower and upper bounds
74 auto it_low = std::lower_bound(values_to_compress.begin(), values_to_compress.end(), current_p_sum - R);
75 auto it_high = std::upper_bound(values_to_compress.begin(), values_to_compress.end(), current_p_sum - L);
76
77 int compressed_low_idx = (it_low == values_to_compress.end()) ? values_to_compress.size() : (it_low - values_to_compress.begin());
78 int compressed_high_idx = (it_high == values_to_compress.begin()) ? -1 : (it_high - values_to_compress.begin() - 1);
79
80 // Count prefix_sums[i-1] in the required range
81 count += ft.query_range(compressed_low_idx, compressed_high_idx);
82
83 // Add current_p_sum to the Fenwick tree
84 ft.update(compressed_coords[current_p_sum], 1);
85 }
86
87 return count;
88}
