Submission #367863

#TimeUsernameProblemLanguageResultExecution timeMemory
367863valerikkSplit the sequence (APIO14_sequence)C++17
100 / 100
1599 ms91440 KiB
#include <bits/stdc++.h>
using namespace std;
typedef long double ld;
typedef long long ll;
const ll INF = (ll)1.1e18;

struct ConvexHull {
    struct Line {
        ll a, b;
        int id;

        ll get(ll x) {
            return a * x + b;
        }

        Line() {}
        Line(ll a, ll b, int id) : a(a), b(b), id(id) {}
    };

    ld inter(Line l1, Line l2) {
        return (ld)(l1.b - l2.b) / (l2.a - l1.a);
    }

    vector<Line> ch;
    vector<ld> interX;
    int pos;

    bool bad(Line l) {
        if ((int)ch.size() <= 1) return false;
        return inter(l, ch[(int)ch.size() - 2]) <= interX.back();
    }

    void addLine(ll a, ll b, int id) {
        if (!ch.empty() && a == ch.back().a && b >= ch.back().b) return;
        Line l = {a, b, id};
        while (bad(l)) ch.pop_back(), interX.pop_back();
        if (ch.empty()) interX.push_back(-INF); else interX.push_back(inter(ch.back(), l));
        ch.push_back(l);
    }

    pair<ll, int> get(ll x) {
        assert(!interX.empty());
        pos = min(pos, (int)ch.size() - 1);
        while (pos + 1 < (int)ch.size() && interX[pos + 1] <= x) pos++;
        return {ch[pos].get(x), ch[pos].id};
    }

    ConvexHull() {
        ch.clear();
        interX.clear();
        pos = 0;
    }
};

ll slow(int n, int k, vector<ll> a, vector<int> &c) {
    vector<ll> pref(n + 1, 0);
    for (int i = 0; i < n; i++) pref[i + 1] = pref[i] + a[i];
    
    auto cost = [&](int l, int r)->ll {
        return (pref[r] - pref[l]) * (pref[r] - pref[l]);
    };

    vector<vector<int>> opt(k + 1, vector<int>(n + 1, -1));
    vector<ll> dp(n + 1, INF);
    for (int i = 1; i <= n; i++) dp[i] = cost(0, i);
    for (int t = 1; t <= k; t++) {
        vector<ll> ndp(n + 1, INF);
        for (int i = 0; i <= n; i++) {
            for (int j = 0; j < i; j++) { 
                if (dp[j] + cost(j, i) < ndp[i]) ndp[i] = dp[j] + cost(j, i), opt[t][i] = j;
            }
        }
        dp.swap(ndp);
    }

    ll res = (cost(0, n) - dp[n]) / 2;

    int cur = n;
    for (int i = k; i >= 0; i--) {
        if (i != k) c.push_back(cur);
        cur = opt[i][cur];
    }
    reverse(c.begin(), c.end());

    return res;
}

ll fast(int n, int k, vector<ll> a, vector<int> &c) {
    vector<ll> pref(n + 1, 0);
    for (int i = 0; i < n; i++) pref[i + 1] = pref[i] + a[i];
    
    vector<vector<int>> opt(k + 1, vector<int>(n + 1, -1));
    vector<ll> dp(n + 1, INF);
    for (int i = 1; i <= n; i++) dp[i] = pref[i] * pref[i];

    for (int t = 1; t <= k; t++) {
        vector<ll> ndp(n + 1, INF);
        ConvexHull ch;
        for (int i = t + 1; i <= n; i++) {
            ch.addLine(-2 * pref[i - 1], pref[i - 1] * pref[i - 1] + dp[i - 1], i - 1);
            auto [res, id] = ch.get(pref[i]);
            ndp[i] = res + pref[i] * pref[i], opt[t][i] = id;
        }
        dp.swap(ndp);
    }

    ll res = (pref[n] * pref[n] - dp[n]) / 2;
    
    int cur = n;
    for (int i = k; i >= 0; i--) {
        if (cur != n) c.push_back(cur);
        cur = opt[i][cur];
    }
    reverse(c.begin(), c.end());

    return res;
}

int main() {
    ios::sync_with_stdio(false); 
    cin.tie(0); 
    int n, k;
    cin >> n >> k;
    vector<ll> a(n);
    for (ll &x : a) cin >> x;
    vector<int> c;
    if (0) cout << slow(n, k, a, c); else cout << fast(n, k, a, c);
    cout << endl;
    for (int i : c) cout << i << " ";
    cout << endl;
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...