Number of partitions

mail@pastecode.io avatar
unknown
c_cpp
a year ago
1.4 kB
12
Indexable
Never
#include<bits/stdc++.h>
using namespace std;

#define ll long long int
#define LD long double

const int N = 100010;

int inf = 1e9;
int mod = 1e9 + 7;

int waysToPartition(vector<int>& nums) {
    int ret = 0, n = nums.size();
    vector<ll> diffs(n, 0);
    unordered_map<ll, int> ls, rs;

    // suffix sum
    for (int i = n - 1; i >= 0; i--)
        diffs[i] = (i == n - 1) ? nums[i] : nums[i] + diffs[i + 1];

    // prefix sum and diffs b/w two partitons divided by pivot i
    for (ll i = 1, ps = nums[0]; i < n; i++) {
        diffs[i] = ps - diffs[i];
        ps = ps + nums[i];
        rs[diffs[i]]++; // keep track of numbers of these diffs
    }

    ret = rs[0]; // doesn't change any number, we need to search the number of 'diff = 0'.

    // change nums[i] to k
    for (int i = 0; i < n; i++) {
        int inc = - nums[i];
        ret = max(ret, (ls.count(inc) ? ls[inc] : 0) + (rs.count(-inc) ? rs[-inc] : 0));
        if (i == n - 1)
            break;
        ls[diffs[i + 1]]++;
        rs[diffs[i + 1]]--;
    }

    return ret;
}

signed main()
{
    //freopen("IN", "r", stdin);
    //freopen("OUT", "w", stdout);

    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);

    int t; cin >> t;

    while (t--) {
        int n;
        cin >> n;

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

        cout << waysToPartition(a) << "\n";
    }
    return 0;
}