Разбор задачи «Разбиение на три равные части»

1. Идея

Если массив можно разрезать на три непустые подряд идущие части с равными суммами, то общая сумма делится на 3, а первая часть заканчивается там, где префиксная сумма равна S/3, вторая — где 2S/3. Задача сводится к подсчёту подходящих пар точек разреза, что делается одним проходом с префиксными суммами.

2. Наблюдения

Пусть S — общая сумма и part = S/3. Нужны две точки разреза: после первой префикс равен part, после второй — 2*part, и вторая строго левее последнего элемента (иначе третья часть пуста). Значит, для каждой позиции, где префикс равен 2*part (и это не самый конец), к ответу прибавляется число более ранних позиций, где префикс равнялся part.

3. Алгоритм

1. Посчитать общую сумму; если она не делится на 3 — ответ 0.
2. Идти по массиву, поддерживая префиксную сумму и счётчик count_part уже встреченных префиксов, равных part.
3. Для позиции j (не последней): если префикс равен 2*part, прибавить к ответу count_part; если префикс равен part, увеличить count_part.
4. Вывести ответ.

4. Почему это работает

Каждая допустимая тройка частей однозначно задаётся двумя точками разреза с префиксами part и 2*part, где первая левее второй, а вторая — не в самом конце. Проходя слева направо и прибавляя число ранее встреченных префиксов part в моменты, когда текущий префикс равен 2*part, мы аккуратно пересчитываем все такие пары ровно один раз.

5. Сложность

Один проход по массиву — O(n) по времени и O(1) дополнительной памяти. Суммы могут быть большими, поэтому используем 64-битные числа, а сам ответ (число способов) тоже может выходить за 32-битный тип.

Код на C++17

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

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

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

    long long one = sum / 3;
    long long two = 2 * one;

    long long pref = 0;
    long long cnt = 0;
    long long ans = 0;

    for (int i = 0; i < n - 1; i++) {
        pref += a[i];
        if (pref == two) {
            ans += cnt;
        }
        if (pref == one) {
            cnt++;
        }
    }

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

Код на Python 3

import sys
data = sys.stdin.buffer.read().split()
n = int(data[0])
a = list(map(int, data[1:1 + n]))

total = sum(a)
if total % 3 != 0:
    print(0)
else:
    part = total // 3
    pref = 0
    count_part = 0
    ways = 0
    for j in range(n - 1):          # j — конец префикса длины j+1, третья часть непуста
        pref += a[j]
        if pref == 2 * part:
            ways += count_part
        if pref == part:
            count_part += 1
    print(ways)