Editorial for count: Chia bánh

Hiểu bài toán

Cho một mảng gồm N phần tử, mỗi phần tử là một số nguyên (có thể âm). Tìm số cách chọn hai chỉ số i, j sao cho 2 <= i <= j <= N-1 và tổng các phần tử từ 1 đến i-1 bằng tổng từ i đến j và bằng tổng từ j+1 đến N. Nói cách khác, chia mảng thành 3 đoạn liên tiếp không rỗng (vì i >= 2 và j <= N-1) sao cho tổng của 3 đoạn bằng nhau. Nếu tổng toàn mảng không chia hết cho 3 thì đáp án là 0. Gọi S là tổng toàn mảng, ta cần tìm số cặp (i, j) sao cho tổng đoạn 1 = S/3, tổng đoạn 1+2 = 2S/3.

Các cách tiếp cận

Cách Brute Force

// O(N^3) hoặc O(N^2)
// Duyệt i từ 2 đến N-2, j từ i đến N-1
// Tính tổng 3 đoạn và so sánh
int count = 0;
for (int i = 2; i <= N - 1; i++) {
    for (int j = i; j <= N - 1; j++) {
        long long sum1 = 0, sum2 = 0, sum3 = 0;
        for (int k = 1; k < i; k++) sum1 += a[k];
        for (int k = i; k <= j; k++) sum2 += a[k];
        for (int k = j + 1; k <= N; k++) sum3 += a[k];
        if (sum1 == sum2 && sum2 == sum3) count++;
    }
}

• Time Complexity: O(N^3) hoặc O(N^2)
• Space Complexity: O(1)

Đây là cách tiếp cận trực tiếp nhất. Duyệt tất cả các cặp (i, j) hợp lệ, tính tổng 3 đoạn và kiểm tra điều kiện. Với N lên tới 500,000, cách này chắc chắn bị Timeout.

Cách Prefix Sum Optimization

#include <iostream>
#include <vector>
using namespace std;

int main() {
    int n;
    cin >> n;
    vector<long long> a(n);
    long long totalSum = 0;
    for (int i = 0; i < n; i++) {
        cin >> a[i];
        totalSum += a[i];
    }

    if (totalSum % 3 != 0) {
        cout << 0 << endl;
        return 0;
    }

    long long target = totalSum / 3;
    long long currentPrefixSum = 0;
    long long countWays = 0;
    long long countPrefixesEqualTarget = 0;

    // Duyệt j từ 0 đến n-2 (đại diện cho kết thúc đoạn 2)
    for (int j = 0; j < n - 1; j++) {
        currentPrefixSum += a[j];

        // Nếu tổng tiền tố hiện tại bằng 2*target, 
        // điều này có nghĩa là đoạn 1+2 có tổng bằng 2S/3.
        // Đoạn 3 tự động có tổng bằng S/3.
        // Ta cần thêm số lượng cách mà đoạn 1 kết thúc ở đây,
        // tức là số lượng chỉ số i trước đó mà tổng tiền tố = target.
        if (currentPrefixSum == 2 * target) {
            countWays += countPrefixesEqualTarget;
        }

        // Nếu tổng tiền tố bằng target, đây là một ứng cử viên cho kết thúc đoạn 1.
        // Điều kiện j <= n-3 đảm bảo rằng còn đủ phần tử cho đoạn 2 và đoạn 3.
        if (currentPrefixSum == target && j <= n - 3) {
            countPrefixesEqualTarget++;
        }
    }

    cout << countWays << endl;
    return 0;
}

• Time Complexity: O(N)
• Space Complexity: O(N) (hoặc O(1) nếu không cần lưu mảng)

Sử dụng tổng tiền tố (Prefix Sum). Nếu tổng toàn mảng là S, ta cần tìm các chỉ số i, j sao cho:

1. Tổng từ 1 đến i-1 là S/3.
2. Tổng từ 1 đến j là 2S/3.

Ta duyệt qua mảng một lần. Vừa duyệt vừa cập nhật tổng tiền tố.

• Nếu tại chỉ số j, tổng tiền tố bằng 2S/3, thì đoạn 1+2 có tổng đúng. Để đoạn 3 cũng đúng, đoạn 1 phải có tổng S/3. Vì ta duyệt từ trái sang phải, ta có thể đếm trước đó có bao nhiêu chỉ số i mà tổng tiền tố tại đó bằng S/3 (và đảm bảo đoạn 2, 3 không rỗng).
• Biến countPrefixesEqualTarget đếm số vị trí có tổng tiền tố bằng S/3.
• Biến countWays cộng dồn kết quả.

Cần chú ý đến độ dài các đoạn (đoạn 1, 2, 3 đều phải có phần tử).

Cách One Pass Counting (Optimized)

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

long long n, a[1000009];

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;
    long long totalSum = 0;
    for (int i = 0; i < n; i++) {
        cin >> a[i];
        totalSum += a[i];
    }

    if (totalSum % 3 != 0) {
        cout << 0 << '\n';
        return 0;
    }

    long long target = totalSum / 3;
    long long countPrefixes = 0;
    long long result = 0;
    long long currentSum = 0;

    // Chỉ duyệt đến n-2 vì j phải <= n-1 (dạng 0-indexed là n-2)
    // để đảm bảo đoạn 3 còn phần tử.
    for (int i = 0; i < n - 1; i++) {
        currentSum += a[i];

        // Nếu currentSum == 2 * target, chúng ta đã tìm thấy một điểm chia thứ 2 (j)
        // hợp lệ. Số cách hình thành là số lượng điểm chia thứ 1 (i) đã xuất hiện trước đó.
        if (currentSum == 2 * target) {
            result += countPrefixes;
        }

        // Nếu currentSum == target, đây là một điểm chia thứ 1 (i) tiềm năng.
        // Chỉ được tính nếu còn đủ chỗ cho đoạn 2 và 3.
        // i <= n-3 nghĩa là sau i còn ít nhất 2 phần tử (một cho đoạn 2, một cho đoạn 3).
        if (currentSum == target && i <= n - 3) {
            countPrefixes++;
        }
    }

    cout << result << '\n';
    return 0;
}

• Time Complexity: O(N)
• Space Complexity: O(N) (để lưu mảng đầu vào)

Đây là phiên bản tối ưu của phương pháp Prefix Sum. Logic: Ta chỉ cần duyệt qua mảng một lần.

• Duy trì biến currentSum là tổng tiền tố hiện tại.
• Duy trì biến countPrefixes để đếm số vị trí i hợp lệ (tổng = target, và đoạn 2, 3 không rỗng).
• Khi currentSum bằng 2 * target, tức là ta đã đi đến một vị trí j hợp lệ cho điểm chia thứ 2. Lúc này, ta cộng countPrefixes vào kết quả.

Lưu ý về điều kiện rỗng: Điểm chia đầu tiên (kết thúc đoạn 1) tại i phải để lại phần tử cho đoạn 2 và 3. Do đó i có thể chạy từ 0 đến n-3 (0-indexed). Điểm chia thứ hai (kết thúc đoạn 2) tại j phải để lại phần tử cho đoạn 3. Do đó j chạy từ 1 đến n-2.

Trong code mẫu: Vòng lặp i chạy từ 0 đến n-2. Nếu currentSum == target và i <= n-3 thì tăng countPrefixes. Nếu currentSum == 2 * target thì cộng countPrefixes vào result. Logic này xử lý đúng các trường hợp tổng âm và các đoạn rỗng.

Phân tích độ phức tạp

Bài học kinh nghiệm

• Tổng toàn mảng S phải chia hết cho 3. Nếu không, đáp án là 0.
• Bài toán quy về tìm 2 chỉ số i, j sao cho tổng tiền tố đến i-1 là S/3 và tổng tiền tố đến j là 2S/3.
• Có thể giải quyết bằng một lần duyệt (O(N)) bằng cách đếm số lượng tiền tố bằng S/3 trước khi gặp tiền tố bằng 2S/3.
• Cần chú ý đến các trường hợp tổng âm và các đoạn rỗng (qi đoạn 1, 2, 3 đều phải chứa phần tử).

Lỗi thường gặp

• Quên kiểm tra chia hết cho 3 của tổng toàn mảng.
• Sai điều kiện biên dẫn đến các đoạn rỗng (ví dụ cho phép i=1 hoặc j=N trong khi các đoạn 1, 2, 3 phải có phần tử).
• Sử dụng biến lưu trữ tổng tiền tố bị tràn số (nên dùng long long vì W_i có thể lên tới 10^9 và N lên tới 5*10^5).
• Xử lý sai trường hợp các đoạn có tổng bằng nhau bằng 0 (ví dụ dãy 0, 0, 0, 0).