Submission #473930

#	Time	Username	Problem	Language	Result	Execution time	Memory
473930		model_code	Diversity (CEOI21_diversity)	C++17	100 / 100	2296 ms	4732 KiB

diversity

#include <cstdio>
#include <cassert>
#include <algorithm>
#include <vector>
#include <cstring>
#include <iostream>
#include <cmath>
#include <string>
#include <deque>

#define FOR(i, a, b) for (int i=(a); i<(b); i++)
#define REP(i, n) FOR(i, 0, n)
#define TRACE(x) cerr << #x << " = " << x << endl
#define _ << " _ " <<
#define X first
#define Y second

using namespace std;

typedef pair<int, int> P;
typedef long long ll;

const int MAX = 300100, SQRT = 550;

int total_freq[MAX];
int freq[MAX];
int no_small[SQRT]; //number of small values with this frequency
vector <int> Large_vals;

inline bool is_small(int val) {
  return total_freq[val] < SQRT;
}

void update(int val, int change) {
  if (is_small(val)) no_small[freq[val]]--;
  freq[val] += change;
  if (is_small(val)) no_small[freq[val]]++;
}

ll sum_squares_arithmetic(ll st, ll a, ll t) { //st^2 + (st+a)^2 + (st+2a)^2 + ... (st+(t-1)a)^2
  return st*st * t + t*(t-1) * a * st + a*a * (((t-1) * t * (2*t-1))/6);
}

int tmp_no_small[SQRT];
vector <P> get_vals() {
  FOR(i, 1, SQRT)
    tmp_no_small[i] = no_small[i];

  for (auto val : Large_vals)
    if (freq[val] < SQRT) tmp_no_small[freq[val]]++;

  vector <P> Vals; //could have repeated values but Vals.size() = O(sqrt(n))
  FOR(i, 1, SQRT) if (tmp_no_small[i]) Vals.push_back(P(i, tmp_no_small[i]));

  sort(Large_vals.begin(), Large_vals.end(), [] (const int &a, const int &b) { return freq[a] < freq[b]; });
  for (auto val : Large_vals)
    if (freq[val] >= SQRT)
      Vals.push_back(P(freq[val], 1));

  return Vals;
}

ll eval(vector <P> Vals) {
  int N = 0, K = 0;
  for (auto it : Vals) {
    N += it.X * it.Y;
    K += it.Y;
  }

  int curr_mod = 1;
  int pref_sum[2] = {0};

  ll ret = 0;
  for (auto it : Vals) {
    assert(it.Y > 0);

    int no_new_odd = (it.Y + curr_mod) / 2;
    int no_new_even = (it.Y + 1-curr_mod) / 2;

    assert(no_new_odd + no_new_even == it.Y);

    //odds
    ret += sum_squares_arithmetic(pref_sum[1], it.X, no_new_odd);
    ret += sum_squares_arithmetic(N - pref_sum[1] - it.X, -it.X, no_new_odd);

    //evens
    ret += sum_squares_arithmetic(pref_sum[0], it.X, no_new_even);
    ret += sum_squares_arithmetic(N - pref_sum[0] - it.X, -it.X, no_new_even);

    pref_sum[0] += it.X * no_new_even;
    pref_sum[1] += it.X * no_new_odd;

    curr_mod = (curr_mod + it.Y) % 2;
  }

  return ((ll)K*N*(N+1) - (ll)N*(K-1) - ret) / 2;
}

int n, q;
int p[MAX];
int qa[MAX], qb[MAX];
ll ans[MAX];
vector <int> Queries[SQRT];

void load() {
  scanf("%d%d", &n, &q);

  REP(i, n) scanf("%d", &p[i]);
  vector <int> V(p, p+n);
  sort(V.begin(), V.end());
  V.resize(unique(V.begin(), V.end()) - V.begin());

  REP(i, n) p[i] = (int) (lower_bound(V.begin(), V.end(), p[i]) - V.begin());

  REP(i, n)
    total_freq[p[i]]++;

  REP(i, n)
    if (total_freq[i] >= SQRT)
      Large_vals.push_back(i);

  REP(i, q) {
    scanf("%d%d", &qa[i], &qb[i]); qa[i]--; qb[i]--;
    Queries[qa[i] / SQRT].push_back(i);
  }

  REP(i, SQRT)
    sort(Queries[i].begin(), Queries[i].end(), [] (const int &a, const int &b) { return qb[a] < qb[b]; });
}

void run(int bucket) {
  memset(freq, 0, sizeof freq);
  memset(no_small, 0, sizeof no_small);

  vector <int> :: iterator itq = Queries[bucket].begin();

  for (int i=bucket*SQRT; i<n; i++) {
    update(p[i], 1);
    for (; itq != Queries[bucket].end() && qb[*itq] == i; itq++) {
      for (int j=bucket*SQRT; j<qa[*itq]; j++)
        update(p[j], -1);

      ans[*itq] = eval(get_vals());

      for (int j=bucket*SQRT; j<qa[*itq]; j++)
        update(p[j], 1);
    }
  }
}

int main()
{
  load();
  REP(i, SQRT) run(i);
  REP(i, q) printf("%lld\n", ans[i]);

  return 0;
}

Compilation message (stderr)

diversity.cpp: In function 'void load()':
diversity.cpp:106:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  106 |   scanf("%d%d", &n, &q);
      |   ~~~~~^~~~~~~~~~~~~~~~
diversity.cpp:108:18: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  108 |   REP(i, n) scanf("%d", &p[i]);
      |             ~~~~~^~~~~~~~~~~~~
diversity.cpp:123:10: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  123 |     scanf("%d%d", &qa[i], &qb[i]); qa[i]--; qb[i]--;
      |     ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~

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