Codeforces Round #629 (Div. 3) F - Make k Equal （离散化 树状数组维护前缀和）

 1 #include<bits/stdc++.h>
 2 using namespace std;
 3 typedef long long ll;
 4 const int maxn = 2e5+5;
 5 ll c[maxn],a[maxn],f[maxn];
 6 int n,k;
 7 struct node{
 8     ll cnt,num;
 9     bool operator < (const node &b){
10         return num<b.num ;
11     }
12 }p[maxn],t[maxn];
13 ll lowbit(ll x){
14     return x& -x;
15 }
16 void add(ll x,ll k){
17     while(x<=n){
18         c[x] = c[x] + k;
19         x = x + lowbit(x);
20     }
21 }
22 ll getsum(ll x){
23     ll ans = 0;
24     while(x>=1){
25         ans +=c[x];
26         x = x-lowbit(x);
27     }
28     return ans;
29 }
30 int main(){
31     scanf("%d%d",&n,&k);
32     for(int i = 1;i<=n;i++) {
33        scanf("%lld",&p[i].num );    
34     }
35     sort(p + 1,p + 1 +n);//离散化
36     t[1].num = p[1].num,t[1].cnt = 1,f[1] = 1; 
37     int cur = 1;
38     for(int i = 2;i<=n;i++){
39         if(p[i].num !=t[cur].num ) t[cur+1].num = p[i].num ,t[cur+1].cnt = 1,cur++,f[cur]=f[cur-1]+1;
40         else t[cur].cnt ++,f[cur]++;//处理f函数和t数组
41     }
42     for(int i = 1;i<=cur;i++){
43         add(i,t[i].cnt *t[i].num );//维护前缀和
44     }
45     ll ans = 1e15+5;
46     for(int i = 1;i<=cur;i++){
47         if(k<=t[i].cnt ) {ans = 0;break;} //存在k>=cnt[i]直接break，ans = 0
48         else{
49             ll d = k - t[i].cnt ;//剩余需要d个移动到num[i]
50             if(f[cur]-f[i]>=d){
51                 ll sum = getsum(cur) - getsum(i) - (f[cur]-f[i]) *(t[i].num+1);
52                 ans = min(ans,sum+d);
53                 
54             }
55             if(f[i-1]>=d){
56                 ll sum = f[i-1]*(t[i].num -1)- getsum(i-1);
57                 ans = min(ans,sum+d);
58             }
59             if(f[cur]-f[i]<d && f[i-1]<d){
60                 ll sum = getsum(cur) - getsum(i) - (f[cur]-f[i]) *(t[i].num+1);
61                 sum = sum + f[i-1]*(t[i].num-1) - getsum(i-1);
62                 ans = min(ans,sum+d);            
63             }
64         }
65     }
66     cout<<ans;
67     return 0;
68 }