Submission

Score: 76

User: kungarooo

Problemset: แคง (Kang)

Language: cpp

Time: 2.092 second

Submitted On: 2025-08-03 20:52:30

#include <bits/stdc++.h>
using namespace std;
long long int sum=0;
unordered_map<long long int,long long int> c;

struct node{
    node* left;
    node* right;
    long long int val,sz=1,pre;
    long long int prior=rand();
    node(){
        left=nullptr;
        right=nullptr;
    }
    node(long long int p){
        left=nullptr;
        right=nullptr;
        val=p;
        pre=val;
    }
};
void update(node* &x){
    if(x==nullptr)return;
    x->sz=1;
    x->pre=x->val;
    if(x->left!=nullptr)x->sz+=x->left->sz,x->pre+=x->left->pre;
    if(x->right!=nullptr)x->sz+=x->right->sz,x->pre+=x->right->pre;
}
void split(node* a,node* &l,node* &r,long long int key){
    if(a==nullptr)l=r=nullptr;
    else if(a->val<=key){
        l=a;
        split(a->right,l->right,r,key);
        update(l);
    }else{
        r=a;
        split(a->left,l,r->left,key);
        update(r);
    }
}
node* merge(node* a,node* b){
    if(a==nullptr||b==nullptr)return (a==nullptr?b:a);
    node* &s=(a->prior > b->prior? a:b),*&t=(a->prior > b->prior? b:a);
    if(s->val < t->val){
        s->right = merge(s->right,t);
        update(s);
        return s;
    }else{
        s->left=merge(s->left,t);
        update(s);
        return s;
    }
    return nullptr;
}
void insert2(node* &a,node* &c){
    if(a==nullptr)a=c;
    else if(a->prior > c->prior){
        insert2((c->val <= a->val ? a->left:a->right),c);
        update(a);
    }else{
        split(a,c->left,c->right,c->val),a=c;
        update(a);
    }
}
void insert(node* &a,long long int c){
    node* x=new node(c);
    insert2(a,x);
}

void del(node* &a,long long int c){
    if(a==nullptr)return;
    if(c==a->val){
        node* x=merge(a->left,a->right);
        //may broke
        //delete th;
        a=x;
        update(a);
    }else if(c < a->val){
        del(a->left,c);
        update(a);
    }else{
        del(a->right,c);
        update(a);
    }
}
long long int find_partition(node* head,int m){
    if(head==nullptr||m<=0)return 0;
    if(head->sz <= m)return head->pre;
    
    if(head->right != nullptr){
        if(head->right->sz >= m)return find_partition(head->right,m);
        else{
            long long int sum=0;
            sum+=head->right->pre;
            sum+=head->val;
            sum+=find_partition(head->left,m-head->right->sz-1);
            return sum;
        }
    }else{
        return head->val+find_partition(head->left,m-1);
    }
}
void print_treap(node* head,int sz){
    if(head==nullptr)return;
    cout<<head->val<<" "<<head->prior<<" "<<head->sz<<" "<<head->pre<<" "<<sz<<"\n";
    print_treap(head->left,sz*2);
    print_treap(head->right,sz*2+1);
}
std::vector<long long> capsize(std::vector<int> A, std::vector<int> B){
    srand(0);
    vector<long long> ret;
    int m=B.size();
    sort(A.begin(),A.end());
    for(int i=0;i<A.size();){
        int idx=i;
        while(++i<A.size()&&A[i]==A[idx]){}
        c[A[idx]]=A[idx]*((long long int)(i-idx));
        sum+=A[idx]*((long long int)(i-idx));
    }
    sum+=B[0];
    c[B[0]]+=B[0];
    node* head=nullptr;
    for(auto i:c){
        insert(head,i.second);
    }
    ret.push_back(sum-find_partition(head,1));
    for(auto i=1;i<m;i++){
        //print_treap(head,1);
        del(head,c[B[i]]);
        c[B[i]]+=B[i];
        sum+=B[i];
        insert(head,c[B[i]]);
        ret.push_back(sum-find_partition(head,i+1));
    }    
    return ret;
}
// int32_t main(){
//     srand(0);
//     int n, m;
//     cin >> n >> m;
//     vector<int> A(n), B(m);
//     for(auto &x: A) cin >> x;
//     for(auto &x: B) cin >> x;
//     vector<long long> ans =  capsize(A,B);
//     for(auto x: ans){
//         cout << x <<" ";
//     }
// }
/*
5 4
6 3 9 6 3
4 6 9 6
*/