Submission #998873

#	Time	Username	Problem	Language	Result	Execution time	Memory
998873		andrei_iorgulescu	Diversity (CEOI21_diversity)	C++14	100 / 100	6957 ms	12920 KiB

diversity

#include <bits/stdc++.h>
#pragma GCC optimize("Ofast")

using namespace std;

using ll = long long;

const int B = 1340;
const int V = 3e5;

int n,a[300005],q;
int fr[300005],aibfr[300005];

void update_fr(int pos,int val)
{
    for (int i = pos; i <= V + 1; i += (i & -i))
        aibfr[i] += val;
}

int query_fr(int pos)
{
    int rr = 0;
    for (int i = pos; i > 0; i -= (i & -i))
        rr += aibfr[i];
    return rr;
}

struct query
{
    int l,r,idx;
};

bool cmp(query A1, query A2)
{
    if (A1.l / B != A2.l / B)
        return A1.l < A2.l;
    else
    {
        if ((A1.l / B) % 2 == 0)
            return A1.r < A2.r;
        return A1.r > A2.r;
    }
}

query queries[50005];
ll rasp[50005];
ll ans1,ans2,ans3;
int v[300005];
int nre;

int aibv[300005];
ll aibvii[300005];

void update_v(int pos,int val)
{
    for (int i = pos; i <= V; i += (i & -i))
        aibv[i] += val;
}

int query_v(int pos)
{
    int rr = 0;
    for (int i = pos; i > 0; i -= (i & -i))
        rr += aibv[i];
    return rr;
}

void update_vii(int pos,ll val)
{
    for (int i = pos; i <= V; i += (i & -i))
        aibvii[i] += val;
}

ll query_vii(int pos)
{
    ll rr = 0;
    for (int i = pos; i > 0; i -= (i & -i))
        rr += aibvii[i];
    return rr;
}

ll sum_vii;

void get_initial_ans()
{
    vector<int> poz_fr;
    for (int i = 1; i <= V; i++)
        poz_fr.push_back(fr[i]);
    sort(poz_fr.begin(),poz_fr.end());
    int l1 = 1,r1 = V;
    for (int i = 0; i < poz_fr.size(); i++)
    {
        if (i % 2 == 0)
            v[l1]  = poz_fr[i],l1++;
        else
            v[r1] = poz_fr[i],r1--;
    }
    for (int i = 1; i <= V; i++)
        ans1 += 1ll * v[i] * (v[i] + 1) / 2;
    ll sumvi = 0,sumivi = 0;
    for (int j = 1; j <= V; j++)
    {
        ans3 += 1ll * j * v[j] * sumvi - 1ll * v[j] * sumivi;
        ans2 += 1ll * sumvi * v[j];
        sumvi += v[j];
        sumivi += 1ll * j * v[j];
    }
    for (int i = 1; i <= V; i++)
        sum_vii += 1ll * v[i] * i;
    for (int i = 1; i <= V; i++)
        update_v(i,v[i]);
    for (int i = 1; i <= V; i++)
        update_vii(i,1ll * v[i] * i);
}

void update_structure(int val,int coef)
{
    //cout << val << ' ' << coef << endl;
    ans1 -= 1ll * fr[val] * (fr[val] + 1) / 2;
    ans1 += 1ll * (fr[val] + coef) * (fr[val] + coef + 1) / 2;
    ans2 -= 1ll * fr[val] * (nre - fr[val]);
    ans2 += 1ll * (fr[val] + coef) * (nre - fr[val]);
    int pos;
    if (coef == 1)
        pos = query_fr(fr[val] + 1);
    else
        pos = query_fr(fr[val]) + 1;
    if (pos % 2 == 1)
        pos = (pos + 1) / 2;
    else
        pos = V - (pos - 1) / 2;
    //cout << ans3 << endl;
    //cout << pos << endl;
    int qvp = query_v(pos);
    ll qviip = query_vii(pos);
    ll delta = 1ll * coef * pos * (qvp - v[pos]) - 1ll * coef * pos * (nre - qvp) + 1ll * coef * (sum_vii - qviip) - 1ll * coef * (qviip - 1ll * v[pos] * pos);
    ans3 += delta;
    int new_val = v[pos] + coef;
    update_v(pos,coef);
    update_vii(pos,coef * pos);
    v[pos] += coef;
    fr[val] += coef;
    //cout << fr[val] << ' ' << fr[val + coef] << endl;
    //cout << 1 << endl;
    update_fr(fr[val] - coef + 1,-1);
    //cout << 3 << endl;
    update_fr(fr[val] + 1,1);
    //cout << 2 << endl;
    nre += coef;
    sum_vii += 1ll * v[pos] * pos;
    sum_vii -= 1ll * (v[pos] - coef) * pos;
}

///fr poate fi intre 0 si V, grija la aib!!

signed main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    cin >> n >> q;
    for (int i = 1; i <= n; i++)
        cin >> a[i],fr[a[i]]++;
    get_initial_ans();
    nre = n;
    for (int i = 1; i <= q; i++)
    {
        cin >> queries[i].l >> queries[i].r;
        queries[i].idx = i;
    }
    sort(queries + 1,queries + q + 1,cmp);
    for (int i = 1; i <= V; i++)
        update_fr(fr[i] + 1,1);
    int l = 1,r = n;
    for (int i = 1; i <= q; i++)
    {
        while (r < queries[i].r)
        {
            r++;
            update_structure(a[r],1);
        }
        while (l > queries[i].l)
        {
            l--;
            update_structure(a[l],1);
        }
        while (r > queries[i].r)
        {
            update_structure(a[r],-1);
            r--;
        }
        while (l < queries[i].l)
        {
            update_structure(a[l],-1);
            l++;
        }
        //cout << ans1 << ' ' << ans2 << ' ' << ans3 << endl;
        rasp[queries[i].idx] = ans1 + ans2 + ans3;
    }
    for (int i = 1; i <= q; i++)
        cout << rasp[i] << '\n';
    return 0;
}

Compilation message (stderr)

diversity.cpp: In function 'void get_initial_ans()':
diversity.cpp:91:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   91 |     for (int i = 0; i < poz_fr.size(); i++)
      |                     ~~^~~~~~~~~~~~~~~
diversity.cpp: In function 'void update_structure(int, int)':
diversity.cpp:138:9: warning: unused variable 'new_val' [-Wunused-variable]
  138 |     int new_val = v[pos] + coef;
      |         ^~~~~~~

• Subtask 1: 4 / 4
• Subtask 2: 10 / 10
• Subtask 3: 8 / 8
• Subtask 4: 16 / 16
• Subtask 5: 26 / 26
• Subtask 6: 36 / 36