Trung bình cộng 2

Nguồn: CLUE OJ Độ khó: Trung bình Tổng tiền tố map

Đề bài

Bài 3 (4 điểm): Trung bình cộng

Cho dãy $A$ gồm $n$ số nguyên $a_1, a_2, \ldots, a_n$ và số nguyên $k$. Hãy đếm xem trong dãy có bao nhiêu dãy con liên tiếp có trung bình cộng bằng $k$.

Dữ liệu: vào từ tệp TBC.INP gồm:

• Dòng 1: gồm 2 số $n$ và $k$ $(1 \le n \le 10^5;\ 1 \le k \le 10^9)$.
• Dòng 2: gồm $n$ số nguyên $a_1, a_2, \ldots, a_n$ $(1 \le a_i \le 10^9)$ cách nhau dấu cách.

Kết quả: Ghi ra tệp TBC.OUT gồm một số là kết quả tìm được.

Giới hạn:

• 50% số test ứng với $n \le 10^4$;
• 50% số test ứng với $n \le 10^5$.

Giải thích ví dụ: với $k = 1$, các đoạn có trung bình bằng 1 là 3 đoạn dài 1 $(1)$, 2 đoạn dài 2 $(1, 1)$ và 1 đoạn dài 3 $(1, 1, 1)$ – tổng cộng $6$ đoạn. Mọi đoạn chứa số $2$ đều có trung bình lớn hơn 1.

Lưu ý: trên CLUE OJ bài này đọc/ghi file TBC.INP / TBC.OUT. Code bên dưới tự kiểm tra: có file thì đọc file, không có thì đọc bàn phím.

Ví dụ – input

4 1
1 1 1 2

Ví dụ – output

6

Lời mở đầu

"Trung bình cộng bằng $k$" nghe có vẻ phải chia, phải lo số thực. Nhưng chỉ cần một mẹo nhỏ: trừ $k$ vào mọi phần tử, bài toán biến thành "đếm đoạn con có tổng bằng 0" – một bài kinh điển giải bằng tổng tiền tố + map trong $O(n \log n)$.

Dành cho học viên · Ưu tiênNộp bài trên hệ thống CLUE OJ

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 OJ

Lộ trình lời giải, bấm vào từng phần để nhảy tới:

Bộ test20 testÝ tưởngtrừ k, tổng 0Cách 1O(n²)Cách 2map

Bộ testBộ test chính thức có gì?

Bộ test có 20 test:

Test	n	k lớn nhất	$a_i$ lớn nhất	Tỉ lệ $a_i = k$	Đáp án
1–10	15 … 20	10	9	11%	0 … 6
11–20	100.000	10	9	10%	0 … 9.751.518

Vài điều đáng chú ý khi soi bộ test:

• Nửa đầu (test 1–10) chỉ có $n \le 20$ – nhỏ hơn nhiều so với $n \le 10^4$ đề ghi. Nửa sau (test 11–20) đều có $n = 10^5$.
• Đề cho $a_i, k$ tới $10^9$ nhưng test thật chỉ có $a_i$ từ 1 đến 9 và $k \le 10$. Code vẫn viết theo giới hạn của đề (dùng map, không dùng mảng đếm theo giá trị).
• Test 1, 5, 17 có $k = 10$ lớn hơn mọi $a_i$ nên đáp án bằng 0; test 4, 12 có $k = 1$ nên chỉ những đoạn toàn số 1 mới thoả.

Đối chiếu giới hạn của đề với test thật và số test đúng của từng cách:

Phần	Đề ghi	Test thật	Cách 1 $O(n^2)$	Cách 2 map
50% đầu	$n \le 10^4$	1–10	10/10	10/10
50% sau	$n \le 10^5$	11–20	0/10	10/10

Đáp án có thể rất lớn

Đáp án lớn nhất trong bộ test là 9.751.518 (test 15, $k = 5$ đúng bằng giá trị giữa của các số 1..9). Trường hợp xấu nhất mọi $a_i = k$ thì mọi đoạn đều thoả, đáp án là $\frac{n(n+1)}{2} \approx 5 \cdot 10^9$ – vượt int. Tổng tiền tố cũng có thể tới $10^{14}$. Cả hai phải dùng long long.

Ý tưởngTrừ k rồi đếm đoạn có tổng 0

Đoạn $a_{i..j}$ (dài $len = j - i + 1$) có trung bình bằng $k$ khi và chỉ khi $$a_i + \ldots + a_j = k \cdot len \iff (a_i - k) + (a_{i+1} - k) + \ldots + (a_j - k) = 0.$$ Đặt $b_t = a_t - k$ và tổng tiền tố $p_j = b_1 + \ldots + b_j$ ($p_0 = 0$). Đoạn $b_{i+1..j}$ có tổng 0 khi và chỉ khi $p_i = p_j$.

Thuật toán

Đáp án = số cặp $(i, j)$ với $0 \le i < j \le n$ và $p_i = p_j$. Duyệt $j$ từ 1 đến $n$, giữ map cnt: cnt[t] = số tiền tố $p_0, \ldots, p_{j-1}$ bằng $t$. Tại mỗi $j$: cộng cnt[p_j] vào đáp án, rồi tăng cnt[p_j]. Ban đầu cnt[0] = 1 (ứng với $p_0$).

Em nhập dãy và $k$ rồi bấm Tiếp để xem từng tổng tiền tố "bắt cặp" với các tiền tố bằng nó:

Dãy a: k =

Ví dụ nhanh:

Cách 1Thử mọi đoạn – 50% điểm

Thử mọi điểm đầu $i$, kéo dài điểm cuối $j$ và cộng dồn tổng $s$. Đoạn thoả khi $s = k \cdot (j - i + 1)$ – so sánh bằng phép nhân để khỏi phải chia (tránh số thực). Bản Python trừ $k$ trước nên chỉ cần so s == 0.

C++▶ Chạy tay từng dòng codeCách 1: thử mọi đoạn, cộng dồn tổng và so với k × độ dài.

Dữ liệu vào (stdin) – sửa thoải mái4 1 1 1 1 2

Dữ liệu mẫu:

ChậmVừaNhanhRất nhanh

Không gian$O(n)$

Thời gian$O(n^2)$

Kết quả chạy thật: bản C++ đúng 10/20 test (quá thời gian ở test 11–20), test đúng chạy lâu nhất 0.00s; bản Python đúng 10/20 test (quá thời gian ở test 11–20), test đúng lâu nhất 0.02s. Ăn trọn nửa đầu; với $n = 10^5$ phải xét khoảng $5 \cdot 10^9$ đoạn nên cả C++ lẫn Python đều quá thời gian.

Cách 2Tổng tiền tố + map – Full điểm

C++▶ Chạy tay từng dòng codeCách 2: trừ k, tổng tiền tố p; mỗi p ghép với các tiền tố trước bằng nó (đếm bằng map).

Dữ liệu vào (stdin) – sửa thoải mái4 1 1 1 1 2

Dữ liệu mẫu:

ChậmVừaNhanhRất nhanh

1. cnt[0] = 1: tiền tố rỗng $p_0 = 0$, để đếm được các đoạn bắt đầu từ vị trí 1 (ví dụ $a_1 = k$ thì $p_1 = 0 = p_0$).
2. dem += cnt[p] trước cnt[p]++: chỉ ghép $p_j$ với các tiền tố đứng trước nó, không tự ghép với chính nó.
3. Không cần lưu mảng $a$: đọc tới đâu xử lý tới đó.
4. Có thể thay map bằng unordered_map cho nhanh hơn chút, nhưng map $O(n \log n)$ đã thừa sức.

Không gian$O(n)$

Thời gian$O(n \log n)$

Kết quả chạy thật: bản C++ đúng 20/20 test, test đúng chạy lâu nhất 0.05s; bản Python đúng 20/20 test, test đúng lâu nhất 0.11s.

Tổng kếtKết quả chạy thật trên bộ test

Cách	Ngôn ngữ	Đúng	Test đúng lâu nhất	Ghi chú
1. Thử mọi đoạn $O(n^2)$	C++	10/20	0.00s	Quá thời gian test 11–20
	Python	10/20	0.02s	quá thời gian test 11–20
2. Trừ k + tổng tiền tố + map	C++	20/20	0.05s	Full điểm
	Python	20/20	0.11s	Full điểm (dict)

Ghi chú của mình: thời gian trong bảng là thời gian CPU đo trên máy của mình (giới hạn 1 giây). Mẹo "trừ $k$ cho mọi phần tử" biến bài trung bình cộng thành bài đếm đoạn tổng 0 – hãy nhớ mẹo này, nó xuất hiện rất thường xuyên.

Rất vui khi anh em nào đã đọc đến những dòng cuối cùng này. Anh em đọc chắc cũng mỏi mắt luôn chứ 😅 vậy anh em nghĩ người viết ra hết đống ở trên còn thế nào nữa 🤪. Vậy nên đừng tiếc gì cho mình một lời bình luận hoặc chia sẻ bài viết để động viên mình nè. Chúc anh em học tốt và gặt hái nhiều thành công 😇