Chia phần (CHIAPHAN)

Nguồn: Đề thi HSG Tin học (bài 1 – 6 điểm) Độ khó: Khá Prefix sum Đếm Nhận xét

Đề bài

Bài 1. Chia phần (6 điểm)

Cho dãy số nguyên $A = (a_1, a_2, \ldots, a_n)$. Hãy đếm số cách chia dãy trên thành 4 dãy con gồm các số liên tiếp của $A$ sao cho tổng các số trong mỗi dãy con đều bằng nhau. Chính xác hơn, mỗi cách chia được mô tả bằng bộ 3 chỉ số $(i, j, k)$: $1 \le i \lt j \lt k \lt n$. Trong đó $(a_1, a_2, \ldots, a_i)$ là dãy 1; $(a_{i+1}, a_{i+2}, \ldots, a_j)$ là dãy 2; $(a_{j+1}, a_{j+2}, \ldots, a_k)$ là dãy 3 và $(a_{k+1}, a_{k+2}, \ldots, a_n)$ là dãy 4.

Hai cách chia khác nhau ứng với hai bộ 3 chỉ số $(i, j, k)$ khác nhau.

Dữ liệu: Vào từ file văn bản chiaphan.inp:

• Dòng đầu tiên ghi số nguyên dương $N$ $(N \le 10^6)$.
• Dòng thứ hai ghi $N$ số nguyên $a_1, a_2, \ldots, a_n$ $(|a_i| \le 10^9;\ i = 1, 2, \ldots, n)$; hai số liên tiếp cách nhau bằng một dấu trống.

Kết quả: Ghi ra file văn bản chiaphan.out một số nguyên là số lượng cách chia tìm được.

Chú ý: Kết quả được chấm trong đó:

• 20% có $N \le 50$.
• 20% có $N \le 500$.
• 30% có $N \le 5000$.
• 30% có $N \le 10^6$.

(Trong ảnh đề tên file ghi lúc là chiapha.inp, lúc là chiphan.out, chắc do gõ nhầm. Code trong bài đọc/ghi chuẩn (bàn phím/màn hình); khi đi thi các bạn thêm freopen đúng tên file mà giám khảo yêu cầu nhé.)

Ví dụ – chiaphan.inp

8
1 1 1 1 1 1 1 1

Ví dụ – chiaphan.out

1

Lời mở đầu

Bữa giờ LeetCode hoài sợ anh em ngán 😅. Nay mình đi comment dạo trong một nhóm lập trình, thấy có bạn hỏi bài này (đề HSG, 6 điểm) cũng hay nên đăng để anh em thử sức nhé. Bài này là một ví dụ rất đẹp cho kỹ thuật mảng cộng dồn (prefix sum) kết hợp với đếm: đi từ vét cạn $O(N^3)$ lên $O(N^2)$ rồi về $O(N)$ chỉ bằng vài nhận xét nhỏ. Hy vọng mang đến được giá trị cho các bạn 😇.

Bạn nào muốn thử nộp bài để hệ thống chấm xem có đúng hay không thì link bài đây nhé:

Link nộp bài

Bấm vào nút bên dưới để nộp bài và được hệ thống CLUE OJ chấm tự động.

Nộp bài trên CLUE OJ

Giải thích ví dụ cho bạn nào chưa hiểu đề: dãy có 8 số 1, tổng là $8$, nên mỗi dãy con phải có tổng $8 / 4 = 2$, tức là mỗi dãy con gồm đúng 2 số 1. Chỉ có một cách duy nhất là cắt sau vị trí $i = 2$, $j = 4$, $k = 6$: $(1, 1)\ |\ (1, 1)\ |\ (1, 1)\ |\ (1, 1)$. Vậy đáp án là 1.

Lưu ý là $a_i$ có thể âm hoặc bằng 0, nên có thể có rất nhiều cách chia. Ví dụ dãy 1 0 1 0 1 1 (tổng $4$, mỗi phần tổng $1$) có tới 4 cách: $i \in \{1, 2\}$, $j \in \{3, 4\}$, $k = 5$.

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

Cách 1Vét cạn (i, j, k) · O(N³) · 40%Cách 2Cố định j, đếm 2 bên · O(N²) · 70%Cách 3Đếm 1 lượt duyệt · O(N) · 100%

Cách 1Vét cạn mọi bộ (i, j, k) + mảng cộng dồn – O(N³)

Làm đúng theo đề: dùng 3 vòng lặp lồng nhau để thử mọi bộ $1 \le i \lt j \lt k \lt n$, rồi tính tổng 4 dãy con xem có bằng nhau không. Nếu mỗi lần lại đi cộng từng phần tử thì tốn thêm $O(N)$ nữa, nên mình chuẩn bị sẵn mảng cộng dồn:

Mảng cộng dồn

$p_0 = 0$, $p_t = a_1 + a_2 + \ldots + a_t$. Khi đó tổng đoạn $(a_{l+1}, \ldots, a_r)$ chính là $p_r - p_l$, tính trong $O(1)$. Vậy 4 dãy con có tổng lần lượt là $p_i$, $\ p_j - p_i$, $\ p_k - p_j$, $\ p_n - p_k$.

Source code để các bạn tham khảo nếu làm theo hướng này:

#include <bits/stdc++.h>
using namespace std;

int n;
long long p[1000005];

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cin >> n;
    for(int i = 1; i <= n; i++){
        long long x;
        cin >> x;
        p[i] = p[i - 1] + x;
    }
    long long cnt = 0;
    for(int i = 1; i < n; i++)
        for(int j = i + 1; j < n; j++)
            for(int k = j + 1; k < n; k++){
                long long s1 = p[i];
                long long s2 = p[j] - p[i];
                long long s3 = p[k] - p[j];
                long long s4 = p[n] - p[k];
                if(s1 == s2 && s2 == s3 && s3 == s4) cnt++;
            }
    cout << cnt;
    return 0;
}

import sys

data = sys.stdin.buffer.read().split()
n = int(data[0])
p = [0] * (n + 1)
for i in range(1, n + 1):
    p[i] = p[i - 1] + int(data[i])

cnt = 0
for i in range(1, n):
    for j in range(i + 1, n):
        for k in range(j + 1, n):
            s1 = p[i]
            s2 = p[j] - p[i]
            s3 = p[k] - p[j]
            s4 = p[n] - p[k]
            if s1 == s2 == s3 == s4:
                cnt += 1
print(cnt)

Đánh giá độ phức tạp của cách làm này:

• Không gian: $O(N)$ cho mảng cộng dồn.
• Thời gian: số bộ $(i, j, k)$ là khoảng $\dfrac{N^3}{6}$. Với $N = 500$ chỉ cỡ $2 \cdot 10^7$ phép kiểm tra, chạy vèo vèo; nhưng $N = 5000$ thì đã lên $2 \cdot 10^{10}$.

Không gian$O(N)$

Thời gian$O(N^3)$

Chỉ được 40% điểm

Cách này qua được 2 subtask đầu ($N \le 500$). Từ $N = 5000$ trở lên là TLE. Ngoài ra tổng có thể lên tới $10^6 \cdot 10^9 = 10^{15}$ nên mảng cộng dồn bắt buộc phải là long long.

Cách 2Nhận xét về tổng + cố định j, đếm hai bên – O(N²)

Gọi $S = p_n$ là tổng cả dãy. Nếu 4 dãy con có tổng bằng nhau thì mỗi dãy phải có tổng đúng bằng $\dfrac{S}{4}$. Đặt $q = \dfrac{S}{4}$, điều kiện "4 tổng bằng nhau" tương đương với:

Nhận xét quan trọng

• Nếu $S$ không chia hết cho 4 thì đáp án là $0$ luôn.
• Ngược lại, $(i, j, k)$ là cách chia hợp lệ khi và chỉ khi: $p_i = q$, $\ p_j = 2q$, $\ p_k = 3q$ (khi đó dãy 4 tự động có tổng $S - 3q = q$).

Tức là ta không cần quan tâm đến "tổng 4 đoạn" nữa, mà chỉ cần quan tâm giá trị mảng cộng dồn tại 3 điểm cắt.

Ví dụ dãy 1 0 1 0 1 1: $S = 4$, $q = 1$, mảng cộng dồn là $p = (1, 1, 2, 2, 3, 4)$.

t	1	2	3	4	5	6
at	1	0	1	0	1	1
pt	1	1	2	2	3	4

p = q → chọn làm ip = 2q → chọn làm jp = 3q → chọn làm k(vị trí n không được chọn)

Có 2 vị trí làm được $i$, 2 vị trí làm được $j$, 1 vị trí làm được $k$, và vị trí nào cũng đúng thứ tự trước–sau nên có $2 \cdot 2 \cdot 1 = 4$ cách.

Từ đó ta có cách $O(N^2)$: cố định điểm cắt giữa $j$ (chỉ xét những $j$ có $p_j = 2q$), đếm số vị trí $i \lt j$ có $p_i = q$ (gọi là trai) và số vị trí $j \lt k \lt n$ có $p_k = 3q$ (gọi là phai). Mỗi $i$ bên trái ghép được với mỗi $k$ bên phải, nên $j$ này đóng góp trai * phai cách (quy tắc nhân).

Source code để các bạn tham khảo nếu làm theo hướng này:

#include <bits/stdc++.h>
using namespace std;

int n;
long long p[1000005];

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cin >> n;
    for(int i = 1; i <= n; i++){
        long long x;
        cin >> x;
        p[i] = p[i - 1] + x;
    }
    long long S = p[n];
    if(S % 4 != 0){
        cout << 0;
        return 0;
    }
    long long q = S / 4, cnt = 0;
    for(int j = 2; j < n - 1; j++){
        if(p[j] != 2 * q) continue;
        long long trai = 0, phai = 0;
        for(int i = 1; i < j; i++)
            if(p[i] == q) trai++;
        for(int k = j + 1; k < n; k++)
            if(p[k] == 3 * q) phai++;
        cnt += trai * phai;
    }
    cout << cnt;
    return 0;
}

import sys

data = sys.stdin.buffer.read().split()
n = int(data[0])
p = [0] * (n + 1)
for i in range(1, n + 1):
    p[i] = p[i - 1] + int(data[i])

S = p[n]
if S % 4 != 0:
    print(0)
else:
    q = S // 4
    cnt = 0
    for j in range(2, n - 1):
        if p[j] != 2 * q:
            continue
        trai = 0
        for i in range(1, j):
            if p[i] == q:
                trai += 1
        phai = 0
        for k in range(j + 1, n):
            if p[k] == 3 * q:
                phai += 1
        cnt += trai * phai
    print(cnt)

Đánh giá độ phức tạp của cách làm này:

• Không gian: $O(N)$.
• Thời gian: với mỗi $j$ ta duyệt lại cả dãy một lần ⇒ $O(N^2)$. Trường hợp xấu nhất là dãy toàn số 0 (khi đó $q = 0$ và mọi $j$ đều có $p_j = 2q$).

Không gian$O(N)$

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

Được 70% điểm

Với $N \le 5000$ thì $N^2 = 2.5 \cdot 10^7$, qua ngon. Nhưng $N = 10^6$ thì $10^{12}$ ⇒ TLE ở subtask cuối.

Cách 3Đếm trong một lượt duyệt – O(N)

Ở cách 2, mỗi lần cố định $j$ ta lại đếm lại từ đầu số $i$ bên trái, rất phí. Các bạn để ý là khi duyệt $t$ từ trái sang phải thì những thứ cần đếm đều chỉ phụ thuộc vào phần bên trái, nên mình duy trì 3 biến:

• cnt1 = số vị trí $i$ đã đi qua có $p_i = q$ (số cách chọn điểm cắt thứ nhất).
• cnt2 = số cặp $(i, j)$ với $i \lt j$ đã đi qua, $p_i = q$, $p_j = 2q$ (số cách chọn 2 điểm cắt đầu).
• ans = số bộ $(i, j, k)$ hợp lệ đã tìm được.

Khi đứng ở vị trí $t$ ($1 \le t \lt n$):

• Nếu $p_t = 3q$: $t$ làm được điểm cắt thứ 3, ghép với mọi cặp $(i, j)$ đứng trước ⇒ ans += cnt2.
• Nếu $p_t = 2q$: $t$ làm được điểm cắt thứ 2, ghép với mọi $i$ đứng trước ⇒ cnt2 += cnt1.
• Nếu $p_t = q$: $t$ làm được điểm cắt thứ nhất ⇒ cnt1++.

Thứ tự 3 lệnh if rất quan trọng!

Khi $S = 0$ thì $q = 2q = 3q = 0$, một vị trí $t$ có thể thỏa cả 3 điều kiện cùng lúc. Vì đề yêu cầu $i \lt j \lt k$ chặt, nên tại mỗi $t$ ta phải cập nhật ans trước (dùng cnt2 của các vị trí trước $t$), rồi mới tới cnt2, cuối cùng mới tới cnt1. Nếu viết ngược thứ tự (cnt1++ trước) thì $t$ sẽ tự ghép với chính nó ⇒ đếm sai. Các bạn thử bấm test "Toàn số 0" ở khung chạy tay bên dưới để thấy nhé.

Vì sao chỉ duyệt t < n?

Điểm cắt $k$ phải nhỏ hơn $n$ để dãy 4 không rỗng. Còn điểm cắt $i \ge 1$ thì đảm bảo dãy 1 không rỗng; dãy 2, dãy 3 không rỗng nhờ $i \lt j \lt k$ đã nói ở trên.

Minh hoạ một lượt duyệt

Các bạn nhập dãy số rồi bấm Tiếp để xem 3 biến cnt1, cnt2, ans thay đổi thế nào khi $t$ đi từ trái sang phải (xanh dương: $p_t = q$, tím: $p_t = 2q$, xanh lá: $p_t = 3q$, viền cam: vị trí đang xét):

Dãy a

Ví dụ nhanh:

Source code để các bạn tham khảo nếu làm theo hướng này:

#include <bits/stdc++.h>
using namespace std;

int n;
long long p[1000005];

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cin >> n;
    for(int i = 1; i <= n; i++){
        long long x;
        cin >> x;
        p[i] = p[i - 1] + x;
    }
    long long S = p[n];
    if(S % 4 != 0){
        cout << 0;
        return 0;
    }
    long long q = S / 4;
    long long cnt1 = 0, cnt2 = 0, ans = 0;
    for(int t = 1; t < n; t++){
        if(p[t] == 3 * q) ans += cnt2;
        if(p[t] == 2 * q) cnt2 += cnt1;
        if(p[t] == q) cnt1++;
    }
    cout << ans;
    return 0;
}

import sys

data = sys.stdin.buffer.read().split()
n = int(data[0])
p = [0] * (n + 1)
for i in range(1, n + 1):
    p[i] = p[i - 1] + int(data[i])

S = p[n]
if S % 4 != 0:
    print(0)
else:
    q = S // 4
    cnt1 = cnt2 = ans = 0
    for t in range(1, n):
        if p[t] == 3 * q:
            ans += cnt2
        if p[t] == 2 * q:
            cnt2 += cnt1
        if p[t] == q:
            cnt1 += 1
    print(ans)

C++▶ Chạy tay từng dòng codeCách 3: duyệt t từ 1 đến n − 1, cập nhật ans, cnt2, cnt1 theo đúng thứ tự.

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

Dữ liệu mẫu:

1#include <bits/stdc++.h>

2using namespace std;

3

4int n;

5long long p[1000005];

6

7int main(){

8 ios_base::sync_with_stdio(false);

9 cin.tie(NULL);

10 cin >> n;

11 for(int i = 1; i <= n; i++){

12 long long x;

13 cin >> x;

14 p[i] = p[i - 1] + x;

15 }

16 long long S = p[n];

17 if(S % 4 != 0){

18 cout << 0;

19 return 0;

20 }

21 long long q = S / 4;

22 long long cnt1 = 0, cnt2 = 0, ans = 0;

23 for(int t = 1; t < n; t++){

24 if(p[t] == 3 * q) ans += cnt2;

25 if(p[t] == 2 * q) cnt2 += cnt1;

26 if(p[t] == q) cnt1++;

27 }

28 cout << ans;

29 return 0;

30}

ChậmVừaNhanhRất nhanh

Đánh giá độ phức tạp của cách làm này:

• Không gian: $O(N)$ cho mảng cộng dồn (thật ra có thể vừa đọc vừa xử lý luôn, nhưng phải biết $S$ trước nên mình vẫn lưu mảng cho dễ hiểu).
• Thời gian: một vòng tính mảng cộng dồn + một vòng duyệt $t$ ⇒ $O(N)$.

Không gian$O(N)$

Thời gian$O(N)$

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

Dãy toàn số 0 với $N = 10^6$ thì mọi bộ $1 \le i \lt j \lt k \le N - 1$ đều hợp lệ ⇒ đáp án là $C^3_{N-1} = 166\,665\,666\,668\,499\,999 \approx 1.7 \cdot 10^{17}$. Vẫn vừa long long, nhưng int thì tràn chắc chắn. Các biến cnt1, cnt2, ans và mảng p đều phải là long long nhé.

Ghi chú của mình: code C++ và Python của cả 3 cách cho đúng output ví dụ của đề, và đã được kiểm tra với 400 bộ test ngẫu nhiên nhỏ (có số âm, nhiều số 0, tổng bằng 0), so với một chương trình vét cạn độc lập cộng thẳng tổng 4 đoạn (riêng bản Python cách 2 thử 150 bộ) – khớp hết. Với $N = 10^6$ (toàn số 0, hoặc $a_i = \pm 10^9$ ngẫu nhiên): cách 3 C++ chạy khoảng 0.15 giây, Python khoảng 1.3 giây. Cách 2 C++ với $N = 5000$ chạy khoảng 0.05 giây, còn $N = 5 \cdot 10^4$ toàn số 0 đã mất khoảng 2 giây nên $N = 10^6$ chắc chắn TLE đúng như phân tích.

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 😇