SEQ
Đề bài
Cho dãy số gồm $n$ số nguyên $a_1, a_2, \ldots, a_n$ và 2 số nguyên không âm $L, R$ $(L \le R)$.
Yêu cầu: Đếm số cặp $(i, j)$ thoả mãn điều kiện: $i \le j$ và $L \le |a_i + \ldots + a_j| \le R$.
Dữ liệu vào:
- Dòng đầu tiên chứa 3 số nguyên $n, L, R$ $(n \le 10^5;\ 0 \le L \le R \le 10^9)$.
- Dòng thứ hai chứa $n$ số nguyên $a_1, a_2, \ldots, a_n$ $(|a_i| \le 10^9)$.
Kết quả: Gồm một số nguyên duy nhất là số lượng cặp $(i, j)$ đếm được.
Hạn chế:
- Có 50% số test ứng với $0 < n \le 10^3$.
- Có 50% số test ứng với $10^3 < n \le 10^5$.
3 0 1
1 -1 24Giải thích ví dụ: với $a = [1, -1, 2]$, $L = 0$, $R = 1$ có 4 cặp thoả mãn: $(1, 1)$ có $|1| = 1$; $(1, 2)$ có $|1 - 1| = 0$; $(2, 2)$ có $|-1| = 1$; $(2, 3)$ có $|-1 + 2| = 1$. Còn $(1, 3)$ có tổng bằng 2 và $(3, 3)$ có tổng bằng 2 đều lớn hơn $R$.
Lời mở đầu
Bài này mình được một em học viên hỏi và mình đã hướng dẫn chi tiết cách thức giải quyết cho em ấy. Tin rằng sẽ có giá trị với nhiều bạn nên mình chia sẻ đề bài để các bạn thử sức làm và kèm thêm đó là toàn bộ tin nhắn của mình với em ấy để bạn nào không giải quyết được thì có thể tham khảo nhen. Hi vọng mang đến được giá trị cho các bạn. Chúc các bạn học tốt và gặt hái nhiều thành công 😇.
Link này chỉ học viên đã tham gia tổ chức mới vào được. Học viên nên ưu tiên nộp ở đây vì sẽ được chấm nhanh hơn.
Nộp bài trên CLUE OJBạn nào muốn thử nộp bài để hệ thống chấm xem có đúng hay không thì dùng link này nhé. Chưa có tài khoản thì tạo tài khoản rồi đăng nhập nộp bài thử hen.
Nộp bài trên CodeforcesLộ trình lời giải, bấm vào từng phần để nhảy tới:
Ý tưởngVì sao phải dùng cấu trúc dữ liệu dạng cây?
Theo anh thì với cái bài này phải dùng đến các cấu trúc dữ liệu dạng cây. Ví dụ trước giờ em có từng dùng map, set bên C++ chưa? Cấu trúc dữ liệu map, set để thống kê đếm tần suất á, thì map, set của C++ được cài đặt theo dạng cây, cụ thể là cây đỏ đen (red black tree). Vì là dạng cây nên các thao tác trên cây như thêm, tìm kiếm, xoá thì chỉ là $O(\log N)$ với $N$ là số lượng phần tử trên cây.
Tại vì nó như sau, em hiểu thế này nè: giả sử nếu cho em 1 dãy số cố định và yêu cầu em thực hiện nhiều lần truy vấn tìm kiếm xem có bao nhiêu giá trị nằm trong đoạn left → right nào đó thì quá đơn giản, em cứ sort dãy số đó là xong, sort 1 lần thôi sau đó cứ làm việc trên nó.
Nhưng nếu như dãy số đó chưa hề cố định mà dãy số đó vẫn đang trong quá trình thêm liên tục vào: ví dụ ngay lúc này dãy số chỉ có 1 phần tử, tiếp theo sẽ thêm 1 phần tử nữa là tổng cộng có 2 phần tử, tiếp theo sẽ thêm 1 phần tử nữa là tổng cộng có 3 phần tử, và cứ thế… tức là nó luôn có sự biến động (thêm hay xoá liên tục). Nhưng mà em phải luôn đảm bảo dù thêm mới hay xoá đi thì dãy số phải luôn có thứ tự tăng dần để từ đó em mới có thể tìm kiếm nhị phân được (tìm kiếm nhị phân điều kiện tiên quyết là dãy số phải có thứ tự chứ dãy mà lộn xộn là sai).
Anh nói ví dụ nếu giả sử dãy đang có 3 phần tử là: 5 7 10 thì đã có thứ tự tăng dần, em tìm kiếm nhị phân được. Nhưng giả sử giờ thêm 1 phần tử vào, giả sử phần tử thêm là số 1 đi thì dãy thành: 5 7 10 1 thì đã không còn thứ tự tăng dần nữa nên không tìm nhị phân được. Nếu muốn tìm thì phải sort lại để thành: 1 5 7 10. Nhưng mà nếu gọi sort liên tục như vậy thì hàm sort độ phức tạp của nó là $O(N \log N)$ với $N$ là số lượng phần tử của danh sách, em cứ tưởng tượng nó sẽ gọi hàm sort $N$ lần thì như thế rất dở, tính tổng lại nó sẽ là $N \log N + (N-1) \log (N-1)$ $+ (N-2) \log (N-2)$ $+ \ldots$ thiếu điều nó còn tệ hơn cả thuật $O(N^2)$, thì chắc chắn là không ổn, sẽ không thể làm thế được.
Vậy nên lúc này cần có những cấu trúc dữ liệu dạng cây như anh đã nói hồi nãy: các thao tác thêm, xoá, tìm kiếm trên cây chỉ có độ phức tạp $O(\log N)$, vì thế ví dụ cần thêm liên tục $N$ lần vào cây thì tổng độ phức tạp sau cùng sẽ là $O(\log N) + O(\log N) + \ldots + O(\log N)$ (cộng $N$ lần liên tục), chẳng phải là $O(N \log N)$ sao? Nên nó sẽ đảm bảo chạy trong 1 giây.
Biến đổiMảng cộng dồn + phá trị tuyệt đối
Thì quay lại bài của em, nếu ta tạo ra mảng prefixSum theo nguyên tắc thì em biết rồi đó: prefixSum[0] = a[0]; rồi với $i$ từ 1 đến $< n$ thì prefixSum[i] = prefixSum[i-1] + a[i]. Nguyên tắc mảng cộng dồn á, này chắc em biết rồi đúng không? Lúc này có thể tính nhanh ra được tổng các phần tử trong đoạn từ index $i$ đến index $j$ sẽ là prefixSum[j] - prefixSum[i-1], tính ra kết quả ngay trong $O(1)$, còn bước tạo ra mảng prefixSum thì tốn $O(N)$ để duyệt qua 1 lần thôi.
prefixSum[i-1] chính là 0 nha.Vậy ta hiểu rằng đề bài nói $L \le |a[i] + \ldots + a[j]| \le R$. Dấu $|\ |$ là trị tuyệt đối, đó chắc em biết trị tuyệt đối là gì đúng không? Nghĩa là nếu bên trong trị tuyệt đối là số âm thì nó biến thành số dương: $|-5| = 5$, còn nếu bên trong là số dương rồi thì thôi nó vẫn vậy. Mà ta nói rằng tổng từ $a[i]$ đến $a[j]$ chính là prefixSum[j] - prefixSum[i-1], nên nghĩa là $L \le |\text{prefixSum}[j] - \text{prefixSum}[i-1]|$ $\le R$. Phá trị tuyệt đối ra đi thì sẽ có 2 tình huống:
- Tình huống 1: $L \le \text{prefixSum}[j] - \text{prefixSum}[i-1]$ $\le R$.
- Tình huống 2: $L \le -(\text{prefixSum}[j] - \text{prefixSum}[i-1])$ $\le R$. Ở tình huống 2 ta phá dấu trừ đi bằng cách đảo lại các cận: $-R \le \text{prefixSum}[j] - \text{prefixSum}[i-1]$ $\le -L$ (em ngẫm xem đúng không? Thử lấy ví dụ đi sẽ thấy rõ nó đúng).
Vậy lúc này cả 2 tình huống thì phần ở giữa đã giống nhau, đều là prefixSum[j] - prefixSum[i-1]. Chuyển vế đổi dấu thì sẽ là:
Như vậy lúc này em xem: $L, R$ đề bài đã cho từ đầu, tức là có dữ liệu. Vậy với mỗi vị trí index $i$ mình xét đến, thì đã có mảng prefixSum rồi, đã có giá trị prefixSum[i-1] rồi nên nghĩa là ta có thể tính ra được các cận. Ở tình huống 1, giả sử L + prefixSum[i-1] gọi là $X$, R + prefixSum[i-1] gọi là $Y$, vậy nghĩa là:
Và đề bài đã ràng buộc là $j \ge i$ đó em, có nghĩa là tại vị trí $j$ từ $i$ hiện tại đang xét cho đến cuối cùng của mảng, đếm xem có bao nhiêu phần tử prefixSum[j] thoả mãn giá trị nằm trong đoạn từ $X$ đến $Y$; đó là với tình huống 1. Còn với tình huống 2 thì cũng y vậy, chỉ là đếm trong đoạn từ -R + prefixSum[i-1] đến -L + prefixSum[i-1] (2 giá trị này tính ra con số rõ ràng được vì mọi dữ liệu đã có).
Đó thì cứ thế với từng vị trí index $i$ xét tới trong mảng prefixSum thì sẽ đếm qua 2 tình huống xem có bao nhiêu vị trí $j$ thoả mãn, thì tổng tất cả các vị trí $j$ thoả mãn sẽ là đáp án sau cùng.
Ví dụChạy tay từ phải qua trái
Giờ hãy nhìn ví dụ sau nhé: nếu ta có mảng có $n = 7$ phần tử là: $1, 3, -5, 8, -10, 2, 6$ và có $L = 5$, $R = 10$, ta sẽ có các cặp $(i, j)$ thoả mãn sau đây với $0 \le i \le j < n$:
| i | j | a[i] + … + a[j] | |sum| |
|---|---|---|---|
| 0 | 3 | 1 + 3 − 5 + 8 | |7| = 7 |
| 0 | 6 | 1 + 3 − 5 + 8 − 10 + 2 + 6 | |5| = 5 |
| 1 | 3 | 3 − 5 + 8 | |6| = 6 |
| 2 | 2 | -5 | |-5| = 5 |
| 2 | 4 | -5 + 8 − 10 | |-7| = 7 |
| 2 | 5 | -5 + 8 − 10 + 2 | |-5| = 5 |
| 3 | 3 | 8 | |8| = 8 |
| 3 | 6 | 8 − 10 + 2 + 6 | |6| = 6 |
| 4 | 4 | -10 | |-10| = 10 |
| 4 | 5 | -10 + 2 | |-8| = 8 |
| 5 | 6 | 2 + 6 | |8| = 8 |
| 6 | 6 | 6 | |6| = 6 |
Tổng cộng có 12 cặp.
Muốn kiểm chứng thì em cứ code cách vét cạn với độ phức tạp $O(N^2)$ để chạy thử ví dụ này của anh là rõ thôi. Thì từ cái mảng $1, 3, -5, 8, -10, 2, 6$ ta có mảng prefixSum như sau:
Giờ em xem nhé, tại vì từ vị trí $i$ thì xét tất cả vị trí $j \ge i$ nên nếu ta xét theo chiều từ trái qua phải, ví dụ khi xét $i = 0$ thì sẽ xét vị trí $j$ là toàn bộ các index $j \ge 0$, nhưng khi xét tiếp vị trí $i = 1$ thì xét vị trí $j$ là toàn bộ $j \ge 1$ (lúc này không tính index 0 nữa) thì như vậy sẽ phải loại trừ ra. Thì thay vì phải loại trừ như vậy ta hãy xét từ phải qua trái đi, tức là xét từ index lớn về nhỏ.
Xét i = 6
Như ví dụ trên $n = 7$ thì index từ 0 đến 6, ta xét index 6 trước, tức là $i = 6$ thì sẽ xét tất cả các index $j \ge 6$, tức là cũng chỉ có mỗi index cuối cùng. Khi $i = 6$ thì prefixSum[i-1] tức là prefixSum[5] = -1.
- Tình huống 1: $5 - 1 \le \text{prefixSum}[j]$ $\le 10 - 1 \iff 4$ $\le \text{prefixSum}[j]$ $\le 9$. Xét trong mảng
prefixSumvới các index $j \ge 6$ xem có bao nhiêu thằng có giá trị nằm trong đoạn từ 4 tới 9? Thì ta thấyprefixSum[6] = 5thoả nằm trong đoạn 4 tới 9 nè, nên xem như tìm thấy 1 phần tử ⇒res = 1, tức là ta có cặp $(i, j) = (6, 6)$ đó em. Với mảng ban đầu thì $a[6] = 6$ thì có phải cũng thoả nằm trong đoạn $L = 5$ đến $R = 10$ không? - Tình huống 2: $-10 - 1$ $\le \text{prefixSum}[j]$ $\le -5 - 1 \iff -11$ $\le \text{prefixSum}[j]$ $\le -6$. Các index $j \ge 6$ chỉ có mỗi index 6 thôi, và tại đó
prefixSum[6] = 5không nằm trong đoạn $-11$ tới $-6$ nên làreskhông thay đổi gì hết và vẫn $= 1$.
Vậy tóm lại sau khi xét vị trí $i = 6$ qua 2 tình huống thì ta chỉ tìm thấy 1 cặp là $(6, 6)$.
Xét i = 5
Rồi tiếp tục xét $i = 5$, khi $i = 5$ thì prefixSum[i-1] tức là prefixSum[4] = -3, thì sẽ xét tất cả các index $j \ge 5$, nghĩa là lúc này xét qua 2 thằng index 5, 6.
- Tình huống 1: $5 - 3 \le \text{prefixSum}[j]$ $\le 10 - 3 \iff 2$ $\le \text{prefixSum}[j]$ $\le 7$. Với $j \ge 5$ chỉ có 2 vị trí $j = 5, 6$ và tại đó
prefixSum[5] = -1,prefixSum[6] = 5. Thì $-1$ không có nằm trong đoạn 2 tới 7, chỉ có 5 nằm trong đoạn thôi nên với tình huống 1 thì ta tìm thấy $j = 6$ thoả mãn, tức là ta có cặp $(i, j) = (5, 6)$. Với mảng ban đầu thì $a[5] + a[6] = 2 + 6 = 8$ có phải cũng thoả nằm trong đoạn $L = 5$ đến $R = 10$ không? ⇒res = 2. - Tình huống 2: $-10 - 3$ $\le \text{prefixSum}[j]$ $\le -5 - 3 \iff -13$ $\le \text{prefixSum}[j]$ $\le -8$. Trong 2 thằng $-1$ và $5$ thì chẳng có thằng nào nằm trong đoạn từ $-13$ tới $-8$ hết nên
reskhông thay đổi gì hết và vẫn $= 2$.
Đó thì quy luật cứ thế chắc em đã rõ. Em bấm Tiếp để xem toàn bộ quá trình đi từ phải qua trái: mỗi bước bỏ prefixSum[i] vào "cây" (luôn được giữ có thứ tự tăng dần), rồi đếm trong cây xem có bao nhiêu giá trị rơi vào 2 đoạn:
Cách 1Vét cạn O(N²) – ăn 50% số test
Trước hết, cách vét cạn rất tự nhiên: thử mọi $i$, rồi cho $j$ chạy từ $i$ tới cuối, vừa chạy vừa cộng dồn sum += a[j] (nên khỏi cần tính lại tổng từ đầu), kiểm tra $L \le |sum| \le R$ thì tăng res. Đây chính là hàm solve1 trong code của anh:
#include <bits/stdc++.h>
using namespace std;
// O(N^2)
long long solve1(long long a[], int n, long long l, long long r){
long long res = 0;
for(int i = 0; i < n; ++i){
long long sum = 0;
for(int j = i; j < n; ++j){
sum += a[j];
if(l <= abs(sum) && abs(sum) <= r){
res++;
}
}
}
return res;
}
int main(){
ios::sync_with_stdio(0);
cin.tie(0);
int n;
long long l, r;
cin >> n >> l >> r;
long long a[n];
for(int i = 0; i < n; ++i){
cin >> a[i];
}
cout << solve1(a, n, l, r);
return 0;
}import sys
def solve1(a, n, l, r):
# O(N^2): thử mọi cặp (i, j), cộng dồn sum khi j tăng
res = 0
for i in range(n):
s = 0
for j in range(i, n):
s += a[j]
if l <= abs(s) <= r:
res += 1
return res
def main():
data = sys.stdin.buffer.read().split()
n, l, r = int(data[0]), int(data[1]), int(data[2])
a = [int(x) for x in data[3:3 + n]]
print(solve1(a, n, l, r))
main()Có $\frac{n(n+1)}{2}$ cặp nên độ phức tạp là $O(N^2)$: với $n \le 10^3$ chỉ khoảng $5 \cdot 10^5$ phép tính, chạy vèo vèo; nhưng với $n = 10^5$ thì là $5 \cdot 10^9$ phép tính, chắc chắn TLE. Lưu ý đáp án có thể lên tới $\frac{n(n+1)}{2} \approx 5 \cdot 10^9$ nên res phải là long long.
Kỹ thuậtOrdered set lưu được số trùng nhau
Nhưng mà vấn đề là ở 2 tình huống, khi xác định được phạm vi cần tìm kiếm xem trong các index $j \ge i$, tức là từ vị trí index $i$ đó đến cuối cùng thì có bao nhiêu phần tử nằm trong đoạn xác định, nếu lúc này mà em dùng 1 vòng for duyệt đếm thì như thế bài của em sẽ là 2 for lồng nhau ⇒ $O(N^2)$. Như anh nói hồi đầu, chỗ này phải dùng cấu trúc cây để lưu trữ: khi đi từ phải qua trái, cứ bỏ các phần tử vào cây, để từ đó khi cần kiểm tra xem trong các phần tử $j \ge i$ có bao nhiêu thằng nằm trong các đoạn của 2 tình huống, tức là kiểm tra TOÀN BỘ CÂY hiện tại đang chứa các phần tử bên trong xem có bao nhiêu thằng nằm trong các đoạn thoả mãn.
Thì lúc này nếu giả sử không phải là cây, mà dùng mảng đi, giả sử khi xét tình huống 1: tính L + prefixSum[i-1] được kết quả, sau đó tìm kiếm nhị phân ta ra được vị trí index $j$ thoả mãn đầu tiên mà tại đó prefixSum[j] >= L + prefixSum[i-1], giả sử là index 5 đi ha; rồi tính R + prefixSum[i-1] được kết quả, sau đó tìm kiếm nhị phân ta ra được vị trí index $j$ thoả mãn cuối cùng mà tại đó prefixSum[j] <= R + prefixSum[i-1], giả sử là index 8 đi ha. Có nghĩa là các index $j$ thoả mãn sẽ là từ 5 đến 8, tức là các index 5, 6, 7, 8 ⇒ có 4 index, có thể tính nhanh bằng cách $8 - 5 + 1 = 4$.
Nhưng vấn đề ở đây nó không phải là mảng, mà nó đang là cây, cây thì nó không có chỉ số index giống như mảng. Ví dụ khi em dùng map để lưu trữ các số, thì khi tìm kiếm nhị phân lower_bound hay upper_bound trên map thì nó chỉ trả về con trỏ đến ngay vị trí node tìm thấy chứ nó không lấy ra được cái chỉ số index như mảng để mà từ đó em lấy 2 cái index trừ nhau để tính ra được số lượng, nó khổ là vậy. Chưa kể 1 điều là map hay set nó cũng chỉ lưu các số không trùng nhau, còn bài này rõ ràng các phần tử có thể trùng nhau thì lại không ổn.
Vậy nên trong tình huống này có 1 dạng cây đặc biệt trong C++ (thư viện __gnu_pbds của g++): bản chất nguyên thuỷ của nó cũng không lưu được số trùng, mà do người ta đã cố tình "độ lại" để phù hợp với các dạng bài cần truy vấn: thay vì lưu giá trị, ta lưu cặp (giá trị, chỉ số), mà chỉ số thì không bao giờ trùng nên cặp cũng không bao giờ trùng. Em hãy xem bài này: How to store duplicate elements in ordered set in C++ – GeeksforGeeks. Nhìn cái code ví dụ của họ nhé, này là thêm dữ liệu vào cây:
int arr[10] = {4, 7, 2, 4, 5, 8, 4, 1, 3, 5};
ordered_set_pair mySet;
for(int i = 0; i < 10; i++){
mySet.insert({arr[i], i});
}Còn lệnh này sẽ đếm và trả ra số lượng các phần tử nhỏ hơn giá trị 5 trong cây:
cout << "Count of elements less than 5: " << mySet.order_of_key({5, -1}) << endl;Nhìn lại cái mảng arr = {4, 7, 2, 4, 5, 8, 4, 1, 3, 5} thì lệnh gọi mySet.order_of_key({5, -1}) nó trả về kết quả là 6, tức là có 6 số trong cây mang giá trị nhỏ hơn 5, ta có thể thấy đó là các số 4, 2, 4, 4, 1, 3. Vì sao lại là -1? Vì mọi cặp $(5, \text{index})$ trong cây đều có index $\ge 0 > -1$ nên đều đứng sau cặp $(5, -1)$, tức là không bị đếm. Thay vì em phải duyệt hết qua mảng với $O(N)$ mới thống kê được thì giờ đây em chỉ cần $O(\log N)$ để duyệt trong cây là ra kết quả.
Đó thì giờ ứng dụng vào, có phải cả 2 tình huống ta đều quy về dạng $X \le \text{prefixSum}[j] \le Y$ đúng không? Thì ta đã có mySet.order_of_key({x, -1}) là trả về số lượng phần tử trong cây nhỏ hơn $x$, thì tư duy 1 chút em sẽ thấy ngay: ta sẽ đếm xem có bao nhiêu thằng nhỏ hơn $Y + 1$ (nhỏ hơn $Y + 1$ tức là $\le Y$ đó), rồi đếm xem có bao nhiêu thằng nhỏ hơn $X$, rồi thì 2 thằng trừ nhau là sẽ ra số lượng thằng nằm trong đoạn $X$ tới $Y$. Em tưởng tượng trục số ra sẽ hiểu chỗ này nhé ^^
Cây đang chứa các giá trị -3, -1, -1, 2, 4, 5, 7, 9. Đếm đoạn $[X, Y] = [0, 5]$: số phần tử nhỏ hơn $Y + 1 = 6$ là 6, số phần tử nhỏ hơn $X = 0$ là 3, vậy có $6 - 3 = 3$ phần tử (chấm xanh) nằm trong đoạn.
Nên lúc này cần đếm xem có bao nhiêu giá trị nằm trong đoạn từ $X$ đến $Y$ chỉ đơn giản là:
mySet.order_of_key({Y + 1, -1}) - mySet.order_of_key({X, -1});① Khi i = 0: em mà gọi prefixSum[i-1] lúc này nó ra index âm là sai đấy nhé. Khi $i = 0$ thì hiểu rằng prefixSum[i-1] chính là 0 nha.
② Khi L = 0: nếu cận $L$ mà bằng 0 thì giữa 2 tình huống sẽ có sự giao nhau. Như cái ví dụ mẫu cho $L = 0, R = 1$ đó, giả sử khi $i = 0$ thì prefixSum[i-1] là 0 như anh đã nói, thì lúc này ở tình huống 1 nó sẽ tìm trong đoạn từ 0 tới 1 và tình huống 2 nó sẽ tìm trong đoạn từ $-1$ tới 0, thì nó đã bị giao nhau ngay chỗ 0 là em sẽ bị tính thêm 1 lần, kết quả là sẽ bị sai. Nên ở tình huống 2 khi xét cận phải đảm bảo nó không bị trùng với tình huống 1 khi $L = 0$; nếu trùng thì phải lùi về 1 đơn vị: tình huống 1 là từ 0 tới 1, tình huống 2 là $-1$ tới 0, thì giờ đây cận 0 đã bị trùng nên tình huống 2 chỉ xét từ $-1$ tới $-1$ thôi. Đó, chỉ có tình huống này là phải lưu ý.
Cách 2Mảng cộng dồn + ordered set – O(N log N)
Rồi giờ anh cho em tham khảo code của anh nhé, có 2 hàm: solve1 là làm theo cách vét cạn, solve2 là làm theo cách tối ưu dùng cấu trúc cây như anh đã nói nãy giờ (hàm main gọi solve2):
#include <bits/stdc++.h>
using namespace std;
// O(N^2)
long long solve1(long long a[], int n, long long l, long long r){
long long res = 0;
for(int i = 0; i < n; ++i){
long long sum = 0;
for(int j = i; j < n; ++j){
sum += a[j];
if(l <= abs(sum) && abs(sum) <= r){
res++;
//cout << i << " " << j << " => |sum| = " << abs(sum) << endl;
}
}
}
return res;
}
// Header files to implement ordered set
// Common file
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
// namespace necessary for GNU based
// policy based data structures
using namespace __gnu_pbds;
// Declaring ordered_set for pair<long long,long long>
typedef tree<pair<long long,long long>, null_type,
less<pair<long long,long long>>, rb_tree_tag,
tree_order_statistics_node_update>
ordered_set_pair;
// O(N * logN)
long long solve2(long long a[], int n, long long l, long long r){
long long prefixSum[n];
prefixSum[0] = a[0];
for(int i = 1; i < n; ++i){
prefixSum[i] = prefixSum[i - 1] + a[i];
}
ordered_set_pair mySet;
long long res = 0;
for(int i = n - 1; i >= 0; --i){
mySet.insert({prefixSum[i], i});
long long x = 0;
if(i > 0){
x = prefixSum[i - 1];
}
res += mySet.order_of_key({r + x + 1, -1}) - mySet.order_of_key({l + x, -1});
res += mySet.order_of_key({-l != l ? -l + x + 1 : -l + x, -1}) - mySet.order_of_key({-r + x, -1});
}
return res;
}
int main(){
ios::sync_with_stdio(0);
cin.tie(0);
int n;
long long l, r;
cin >> n >> l >> r;
long long a[n];
for(int i = 0; i < n; ++i){
cin >> a[i];
}
cout << solve2(a, n, l, r);
return 0;
}import sys
from bisect import bisect_left, bisect_right
def solve2(a, n, l, r):
# mảng cộng dồn
pre = [0] * n
s = 0
for i in range(n):
s += a[i]
pre[i] = s
# Python không có ordered set nên mình dùng cây Fenwick (BIT)
# trên các giá trị prefixSum đã nén (sắp xếp + bỏ trùng)
vals = sorted(set(pre))
m = len(vals)
bit = [0] * (m + 1)
def add(p): # thêm 1 phần tử ở vị trí nén p (1-based)
while p <= m:
bit[p] += 1
p += p & -p
def cnt(p): # số phần tử có vị trí nén <= p
c = 0
while p > 0:
c += bit[p]
p -= p & -p
return c
def count(X, Y): # số phần tử trong cây có giá trị thuộc [X, Y]
if X > Y:
return 0
return cnt(bisect_right(vals, Y)) - cnt(bisect_left(vals, X))
res = 0
for i in range(n - 1, -1, -1):
add(bisect_left(vals, pre[i]) + 1)
x = pre[i - 1] if i > 0 else 0
res += count(l + x, r + x) # tình huống 1
res += count(-r + x, (-l if l != 0 else -1) + x) # tình huống 2 (L = 0 thì bỏ số 0)
return res
def main():
data = sys.stdin.buffer.read().split()
n, l, r = int(data[0]), int(data[1]), int(data[2])
a = [int(x) for x in data[3:3 + n]]
print(solve2(a, n, l, r))
main()Giải thích một số chỗ trong code
mySet.insert({prefixSum[i], i});đi từ phải qua trái nên lúc xét $i$, cây chứa đúng cácprefixSum[j]với $j \ge i$. Lưu kèm chỉ sốiđể các giá trị trùng nhau vẫn là các cặp khác nhau.mySet.order_of_key({r + x + 1, -1}) - mySet.order_of_key({l + x, -1}): tình huống 1, đếm các giá trị trong đoạn $[L + x, R + x]$ với $x$ =prefixSum[i-1].mySet.order_of_key({-l != l ? -l + x + 1 : -l + x, -1}) - mySet.order_of_key({-r + x, -1}): tình huống 2, đếm trong đoạn $[-R + x, -L + x]$. Biểu thức-l != lchỉ sai khi $L = 0$; lúc đó cận phải là $-L + x = x$ không được cộng 1 nữa, tức là chỉ đếm tới $x - 1$, đúng như lưu ý ② ở trên.long long prefixSum[n];vàlong long a[n];là mảng có kích thước thay đổi (VLA), g++ chấp nhận; $|a_i| \le 10^9$ và $n \le 10^5$ nên tổng có thể tới $10^{14}$, bắt buộc phải dùnglong long.
Mỗi vòng lặp có 1 lần insert và 4 lần order_of_key, mỗi thao tác $O(\log N)$, nên tổng cộng $O(N \log N)$.
prefixSum đã nén (sắp xếp + bỏ trùng), đếm số phần tử trong đoạn $[X, Y]$ bằng bisect + 2 lần truy vấn tiền tố, cùng ý tưởng đi từ phải qua trái. Tất cả các bản đã được kiểm tra với 400 bộ test ngẫu nhiên (có cả $L = 0$, $L = R = 0$, số âm, số trùng nhau) so với code vét cạn, đều khớp. Với $n = 10^5$: bản C++ tối ưu chạy khoảng 0.15–0.26 giây, bản Python BIT khoảng 1.5 giây; còn bản vét cạn C++ mất tới 7–10 giây (TLE đúng như dự đoán).