Educational Codeforces Round 69 (Rated for Div. 2) C. Array Splitting

• Time : O(nlogn)
• Space : O(n)

#include <bits/stdc++.h>

#pragma optimization_level 3
#pragma GCC optimize("Ofast,no-stack-protector,unroll-loops,fast-math,O3")
#pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx")
#pragma GCC optimize("Ofast")//Comment optimisations for interactive problems (use endl)
#pragma GCC target("avx,avx2,fma")
#pragma GCC optimization ("unroll-loops")

using namespace std;

struct PairHash {
    inline std::size_t operator()(const std::pair<int, int> &v) const { return v.first * 31 + v.second; }
};

// speed
#define Code ios_base::sync_with_stdio(false);
#define By ios::sync_with_stdio(0);
#define Sumfi cout.tie(NULL);

// alias
using ll = long long;
using ld = long double;

// constants
const ld PI = 3.14159265358979323846;  /* pi */
const ll INF = 1e18;
const ld EPS = 1e-9;
const ll MAX_N = 101010;
const ll mod = 998244353;

// typedef
typedef pair<ll, ll> pll;
typedef vector<pll> vpll;
typedef array<ll,3> all3;
typedef array<ll,5> all5;
typedef vector<all3> vall3;
typedef vector<all5> vall5;
typedef vector<ll> vll;
typedef vector<vll> vvll;
typedef vector<int> vi;
typedef vector<string> vs;
typedef vector<vs> vvs;
typedef unordered_set<ll> usll;
typedef unordered_set<pll, PairHash> uspll;
typedef unordered_map<ll, ll> umll;

// macros
#define rep(i,n) for(ll i=0;i<n;i++)
#define rrep(i,m,n) for(ll i=n;i>=m;i--)
#define all(a) begin(a), end(a)
#define rall(a) rbegin(a), rend(a)

ll __gcd(ll x, ll y) { return !y ? x : __gcd(y, x % y); }
ll __lcm(ll x, ll y) { return x * y / __gcd(x,y); }

using namespace std;

ll solve(vll& A, ll k) {
    if(k == 1) return A.back() - A[0];
    vpll diff;

    ll n = A.size(), res = 0, now = 0;
    rep(i, n - 1) {
        diff.push_back({A[i + 1] - A[i], i});
    }

    sort(all(diff));
    vll cut;
    while(--k) {
        cut.push_back(diff.back().second); diff.pop_back();
    }
    sort(rall(cut));
    while(!cut.empty()) {
        ll e = cut.back(); cut.pop_back();
        res += A[e] - A[now];
        now = e + 1;
    }
    res += A[n - 1] - A[now];

    return res;
}

int main() {
    Code By Sumfi

    //cout.precision(12);

    ll tc = 1;
    //cin>>tc;
    for (ll i = 1; i <= tc; i++) {
        ll n, k;
        cin>>n>>k;
        vll A(n);
        rep(i,n) cin>>A[i];
        cout<<solve(A,k)<<endl;
    }
    return 0;
}