Yet Another Array Partitioning Task [Codeforces Round #538 (Div. 2) B]

https://codeforces.com/contest/1114/problem/B

N要素の配列Aがある。
この配列を「連続していてM要素以上の列」にK個に分ける。
分割後の各列について、大きい方からM個取ってくる。
取ったMK個の総和の最大値を求めて、分割地点を答えよ。

2≦N≦2*10^5
1≦M
2≦K
MK≦N

解説

https://codeforces.com/contest/1114/submission/49712803

実は全体の大きい方からMK個を全て取ることができる。
「取りたい要素MK個を決めた後に、それを取るようにしてK個に分割することは、常に可能である。」
なので、{A[i],i}の配列を降順ソートした順番で取る要素を決定すればいい。
取る要素のA[i]の総和がとりあえず1つ目の答え。
　
2つ目の分割地点は、「use[i] := A[i]をとるか」を作り、頭からuse[i]=1をM個毎に取っていくように区切ればいい。
use配列作成後は、use[i]=1ならcnt++をして、cnt%M=0になったら、M個取ったことになるので、その地点を答えとする。

int N, M, K, A[201010];
int use[201010];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> M >> K;
    rep(i, 0, N) cin >> A[i];

    vector<pair<ll, int>> v;
    rep(i, 0, N) v.push_back({ A[i], i });
    sort(all(v), greater<pair<ll, int>>());

    ll ans = 0;
    rep(i, 0, M*K) {
        ans += v[i].first;
        use[v[i].second] = 1;
    }

    vector<int> ans2;

    int cnt = 0;
    rep(i, 0, N) if(use[i]) {
        cnt++;
        if (cnt % M == 0) ans2.push_back(i + 1);
    }

    printf("%lld\n", ans);
    rep(k, 0, K - 1) {
        if (k) printf(" ");
        printf("%d", ans2[k]);
    }
    printf("\n");
}