Codeforces 21C: Stripe 2 - 前綴和 + 計數構造題解析

題目大意

給長度 n 的整數數列 a[1..n]，將其劃分為三段非空連續段，使三段和相等。求劃分方案數。

核心觀察

• 總和 sum 必須能被 3 整除，否則答案 0。
• 目標段和 target = sum / 3。
• 劃分點 i, j (1 ≤ i < j < n) 滿足：
  prefix[i] == target 且 suffix[j] == target（即 prefix[j-1] == 2*target）。

算法：前綴和 + 後綴計數

1. 計算前綴和 lsum[i] 和後綴和 rsum[i]。
2. 標記 lok[i] = 1 若 lsum[i] == target (i 為第一段末尾)。
3. 標記 rok[i] = 1 若 rsum[i] == target (i 為第三段開頭)。
4. 對 rok 做後綴和：rok[i] += rok[i+1]，表示從 i 到 n-1 有多少合法第三段起點。
5. 遍歷第一段末尾 i = 1 .. n-2，答案累加 lok[i] * rok[i+2]（中間段至少一個元素，故第三段起點 ≥ i+2）。

特殊情況：全 0 數列

若所有 a[i] == 0，任意劃分均合法，方案數為 C(n-1, 2) = (n-1)*(n-2)/2。

完整代碼

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const int N = 100050;
ll a[N], lsum[N], rsum[N];
int lok[N], rok[N];

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n;
    while (cin >> n) {
        ll sum = 0;
        bool allZero = true;
        for (int i = 1; i <= n; ++i) {
            cin >> a[i];
            sum += a[i];
            if (a[i]) allZero = false;
        }
        if (allZero) {
            cout << 1LL * (n - 1) * (n - 2) / 2 << '\n';
            continue;
        }
        if (sum % 3) { cout << "0\n"; continue; }
        ll target = sum / 3;
        for (int i = 1; i <= n; ++i) lsum[i] = lsum[i-1] + a[i];
        for (int i = n; i >= 1; --i) rsum[i] = rsum[i+1] + a[i];
        fill(lok, lok + n + 1, 0);
        fill(rok, rok + n + 1, 0);
        for (int i = 1; i <= n; ++i) if (lsum[i] == target) lok[i] = 1;
        for (int i = n; i >= 1; --i) {
            if (rsum[i] == target) rok[i] = 1;
            rok[i] += rok[i+1];
        }
        ll ans = 0;
        for (int i = 1; i + 2 <= n; ++i) ans += 1LL * lok[i] * rok[i+2];
        cout << ans << '\n';
    }
    return 0;
}

複雜度

• 時間：O(n)，單次掃描。
• 空間：O(n)。

關鍵細節

• lok[n] 不能為 1（第一段不能取完整數列），代碼中循環 i+2<=n 隱含處理。
• 使用 64 位整數 (long long) 防止答案溢出（n ≤ 1e5 時方案數 ~ 1e10）。
• 全 0 情況單獨處理，避免除以 0 或邏輯錯誤。