Codeforces
CF Step
Youtube Linkedin Discord Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Back to homepage

Code : Array Splitting

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

void solve(vector<int> &a, int k) {
    int n = a.size();

    vector<long long> suff(n);
    suff[n - 1] = a[n - 1];
    for (int i = n - 2; i >= 0; i--) {
        suff[i] = suff[i + 1] + a[i];
    }

    // sum(f(i)) : Sum of k elements of suffix array, but 1st element is
    // compulsory
    long long res = suff.front();
    suff.erase(suff.begin());

    sort(suff.begin(), suff.end(), greater<long long>());
    res += accumulate(suff.begin(), suff.begin() + k - 1, 0LL);
    cout << res << "\n";
}

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

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