SRM481-D1-500 BatchSystemRoulette

#define ll long long
#include <bits/stdc++.h>
using namespace std;

/*
Ordered set usage:
order_of_key (k) : Number of items strictly smaller than k.
find_by_order(k) : K-th element in a set (counting from zero).
*/

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
#define ordered_set tree<int, null_type,less<int>, rb_tree_tag,tree_order_statistics_node_update>

const int OO = 1e9;
const double EPS = 1e-9;

class BatchSystemRoulette {
public:
    pair<double,int> arr[50];
    int arr_sz;

    double fpow(double base, int p) {
        if(p == 0) {
            return 1.0;
        }
        if(p&1) {
            return base * fpow(base,p-1);
        }
        double ret = fpow(base,p/2);
        return ret*ret;
    }

    vector <double> expectedFinishTimes(vector <int> duration, vector <string> user) {
        map<string,vector<pair<double,int>>> tasks;
        map<string,ll> sum;
        vector<pair<ll,string>> v;
        map<ll,int> cnt;
        for(int i = 0; i < user.size(); i++) {
            tasks[user[i]].push_back({duration[i]*1.0,i});
            sum[user[i]] += duration[i];
        }
        for(pair<string,ll> p : sum ) {
            cnt[p.second]++;
            v.push_back({p.second,p.first});
        }
        sort(v.begin(),v.end());
        ll pre = 0;
        vector<double> ans;
        ans.resize(duration.size());
        for(int i = 0; i < v.size(); i++) {
            pair<ll,string> us = v[i];
            arr_sz = 0;
            double eq_wait = 0;
            if(cnt[us.first] > 1) {
                eq_wait = us.first * fpow(2,cnt[us.first]-1) + (cnt[us.first]-1) * us.first * fpow(2, cnt[us.first]-2);
                eq_wait /= fpow(2, cnt[us.first]-1);
                eq_wait -= us.first;
            }
            for(pair<double,int> p : tasks[us.second]) {
                arr[arr_sz] = p;
                arr_sz++;
            }
            for(int j = 0; j < arr_sz; j++) {
                double curr = 0;
                for(int k = 0; k < arr_sz; k++) {
                    if(k == j)
                        curr += arr[k].first * fpow(2, arr_sz-1);
                    else
                        curr += arr[k].first * fpow(2, arr_sz-2);
                }
                ans[arr[j].second] = curr / fpow(2,arr_sz-1) + pre + eq_wait;
            }
            if(i < v.size()-1 && v[i+1].first != us.first)
                pre += us.first * cnt[us.first];
        }
        return ans;
    }
};

int main()
{
    /*ios_base::sync_with_stdio(NULL);
    cin.tie(0);*/
    BatchSystemRoulette bs;
    vector<double> v = bs.expectedFinishTimes({14,13,14,13},{"A", "A", "B", "B"});
    for(int i = 0; i < v.size(); i++) {
        cout << v[i] << " ";
    }
    return 0;
}