This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include<bits/stdc++.h>
using namespace std;
int main(){
ios_base::sync_with_stdio(false);
cin.tie(0);
int n,k;
cin>>n>>k;
vector<long long>arr(n);
for (int i = 0;i<n;++i)cin>>arr[i];
vector<long long>pref(n + 1,0);
for (int i = 0;i<n;++i){
pref[i + 1] = pref[i] + arr[i];
}
auto getscore = [&](int l,int r){
return pref[r] - pref[l];
};
set<int>index;
vector<int>ans(k);
set<pair<long long,long long>>brr;
long long score = 0;
long long answer=0;
for (int i = 0;i<k;++i){
brr.insert({arr[i],i});
ans[i] = i + 1;
score+=getscore(i,i + 1) * getscore(i + 1,n);
index.insert(i);
}
vector<int>cur = ans;
answer = score;
for (int i = k;i<n - 1;++i){
ans.push_back(i + 1);
auto u = *brr.begin();
score+=getscore(i,i + 1) * getscore(i + 1,n);
index.insert(i);
brr.insert({arr[i],i});
brr.erase(u);
ans[u.second] = -1;
auto it = index.find(u.second);
bool ok1 = true,ok2 = true;
if (next(it)==index.end())ok1 = false;
if (it == index.begin())ok2 = false;
int nxt = *next(it);
int p = *prev(it);
index.erase(it);
int score1 = (ok2)?(getscore(u.second + 1,n) * (arr[p] + arr[u.second]) - getscore(p + 1,n) * arr[p] - getscore(u.second + 1,n) * arr[u.second]):0;
int score2 = (ok1)?(getscore(nxt + 1,n) * (arr[nxt] + arr[u.second]) - getscore(nxt + 1,n) * arr[nxt] - getscore(u.second + 1,n) * arr[u.second]):0;
if ((ok1 && score2 > score1) || !ok2){
auto v = make_pair(arr[nxt],nxt);
brr.erase(v);
score+=score2;
arr[nxt]+=arr[u.second];
brr.insert({arr[nxt],nxt});
}
else if (ok2){
auto v = make_pair(arr[p],p);
brr.erase(v);
score+=score1;
arr[p]+=arr[u.second];
brr.insert({arr[p],p});
}
if (score > answer){
cur = ans;
}
answer = max(answer,score);
}
cout<<answer<<'\n';
for (auto x:cur){
if (x!=-1)cout<<x<<' ';
}
cout<<'\n';
return 0;
}
/*
7 3
4 1 3 4 0 2 3
*/
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |