Submission #867890

#	Time	Username	Problem	Language	Result	Execution time	Memory
867890		evenvalue	Diversity (CEOI21_diversity)	C++17	64 / 100	7054 ms	13728 KiB

diversity

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

#ifdef evenvalue
  #include "debug.h"
  #define debug(...) print(#__VA_ARGS__, __VA_ARGS__)
#else
  #define debug(...)
#endif

using int64 = long long;
using ld = long double;

template<typename T>
using min_heap = priority_queue<T, vector<T>, greater<T>>;
template<typename T>
using max_heap = priority_queue<T, vector<T>, less<T>>;

namespace Read {
int Int() {
  int x;
  cin >> x;
  return x;
}
int64 Int64() {
  int64 x;
  cin >> x;
  return x;
}
char Char() {
  char c;
  cin >> c;
  return c;
}
string String() {
  string s;
  cin >> s;
  return s;
}
double Double() {
  return stod(String());
}
ld LongDouble() {
  return stold(String());
}
template<typename T1, typename T2>
pair<T1, T2> Pair() {
  pair<T1, T2> p;
  cin >> p.first >> p.second;
  return p;
}
template<typename T>
vector<T> Vec(const int n) {
  vector<T> v(n);
  for (T &x : v) {
    cin >> x;
  }
  return v;
}
template<typename T>
vector<vector<T>> VecVec(const int n, const int m) {
  vector<vector<T>> v(n);
  for (vector<T> &vec : v) {
    vec = Vec<T>(m);
  }
  return v;
}
}// namespace Read

constexpr int kInf = 1e9 + 10;
constexpr int64 kInf64 = 1e15 + 10;
constexpr int kMod = 1e9 + 7;
constexpr int kMaxN = 2e5 + 10;

vector<int> precalc(const int n) {
  vector<int> index(n);

  deque<int> d;
  for (int i = n - 1; i >= 0; i--) {
    (i & 1) ? d.push_front(i) : d.push_back(i);
  }

  for (int i = 0; const int x : d) {
    index[x] = i++;
  }

  return index;
}

template<class node, class F = std::function<node(const node &, const node &)>>
class SegTree {
  int n = 0;
  std::vector<node> t;
  F unite;

  void build(const int x, const int l, const int r, const std::vector<node> &a) {
    if (l == r) {
      t[x] = a[l];
      return;
    }
    const int mid = (l + r) / 2;
    const int y = 2 * (mid - l + 1) + x;
    build(x + 1, l, mid, a);
    build(y, mid + 1, r, a);
    t[x] = unite(t[x + 1], t[y]);
  }

  void update(const int x, const int l, const int r, const int p, const node &v) {
    if (l == p and p == r) {
      t[x] = v;
      return;
    }
    const int mid = (l + r) / 2;
    const int y = 2 * (mid - l + 1) + x;
    if (p <= mid) {
      update(x + 1, l, mid, p, v);
    } else {
      update(y, mid + 1, r, p, v);
    }
    t[x] = unite(t[x + 1], t[y]);
  }

  node query(const int x, const int l, const int r, const int ql, const int qr) const {
    if (ql <= l and r <= qr) {
      return t[x];
    }
    const int mid = (l + r) / 2;
    const int y = 2 * (mid - l + 1) + x;
    if (qr <= mid) {
      return query(x + 1, l, mid, ql, qr);
    } else if (mid < ql) {
      return query(y, mid + 1, r, ql, qr);
    } else {
      return unite(query(x + 1, l, mid, ql, qr), query(y, mid + 1, r, ql, qr));
    }
  }

  void debug_node(const int x, const vector<int> &path) const {
    for (int i = 0; i < path.size(); i++) {
      cerr << path[i] << (i == path.size() - 1 ? ": " : " -> ");
    }
    cerr << t[x] << '\n';
  }

  void print_tree(const int x, const int l, const int r, vector<int> path) const {
    path.push_back(x);
    if (l == r) {
      debug_node(x, path);
      return;
    }
    const int mid = (l + r) / 2;
    const int y = 2 * (mid - l + 1) + x;
    print_tree(x + 1, l, mid, path);
    print_tree(y, mid + 1, r, path);
    debug_node(x, path);
  }

public:
  SegTree() = default;
  explicit SegTree(const int n, const node e, F f) : n(n), t(2 * n - 1, e), unite(std::move(f)) {}
  explicit SegTree(const std::vector<node> &a, F f) : n(a.size()), t(2 * (a.size()) - 1), unite(std::move(f)) {
    build(0, 0, n - 1, a);
  }

  void update(const int p, const node &v) {
    assert(0 <= p and p < n);
    update(0, 0, n - 1, p, v);
  }

  [[nodiscard]] node query(const int l, const int r) const {
    assert(0 <= l and l <= r and r < n);
    return query(0, 0, n - 1, l, r);
  }

  void print_tree() const {
    print_tree(0, 0, n - 1, {});
    cerr << "----------\n\n";
  }
};

struct node {
  int64 val = 0;
  int64 seq = 0;
};

template<typename T>
class CoordinateCompression {
  bool sorted = false;
  vector<T> v;

public:
  void add(const T &x) {
    v.push_back(x);
  }

  void compress() {
    sort(v.begin(), v.end());
    v.erase(unique(v.begin(), v.end()), v.end());
    sorted = true;
  }

  int get(const int x) {
    if (not sorted) assert(false);
    return distance(v.begin(), lower_bound(v.begin(), v.end(), x));
  }

  int size() {
    return v.size();
  }
};

class mo_s {
  struct query {
    int l, r, v, t;
    int64 ans;

    query(const int l, const int r, const int v, const int t) : l(l), r(r), v(v), t(t), ans(0) {}
  };

  const int kBlockSize = 350;

  const int n;
  int kDistinct;
  vector<int> a;
  vector<query> queries;

  [[nodiscard]] int block_of(const int i) const {
    return i / kBlockSize;
  }

public:
  explicit mo_s(vector<int> v) : n(v.size()), a(v) {
    CoordinateCompression<int> cc;
    for (const int x : a) {
      cc.add(x);
    }
    cc.compress();
    for (int &x : a) {
      x = cc.get(x);
    }
    kDistinct = cc.size();
  }

  void add_query(const int l, const int r, const int v, const int t) {
    queries.emplace_back(l, r, v, t);
  }

  vector<int64> solve() {
    sort(queries.begin(), queries.end(), [&](const query &q1, const query &q2) {
      if (block_of(q1.l) != block_of(q2.l)) {
        return (q1.l < q2.l);
      }
      return (block_of(q1.l) & 1 ? q1.r < q2.r : q1.r > q2.r);
    });

    //todo data structure
    const vector<int> index = precalc(kDistinct);

    vector<int> table(kDistinct);
    vector<int> freq(kDistinct);

    SegTree<node> st(kDistinct, node(), [&](const node l, const node r) {
      node ans;
      ans.val = l.val + r.val;
      ans.seq = l.seq + r.seq;
      return ans;
    });

    int64 diversity = 0;

    auto add = [&](const int i) {
      //todo: debug
//      cout << "Called add @" << i << '\n';

      const int val = a[i];
      const int f = freq[val] + 1;
      const int tpos = distance(table.begin(), prev(upper_bound(table.begin(), table.end(), f - 1)));
      freq[val]++;
      table[tpos]++;
      const int rank = index[tpos];
      st.update(rank, {f, rank * 1LL * f});

      const node left = st.query(0, rank);
      const node right = st.query(rank, kDistinct - 1);

      diversity += (rank * left.val) - left.seq;
      diversity += right.seq - (rank * right.val);

//      cout << "After adding, diversity = " << diversity << '\n';
    };

    auto rem = [&](const int i) {
      //todo: debug
//      cout << "Called rem @" << i << '\n';

      const int val = a[i];
      const int f = freq[val] - 1;
      const int tpos = distance(table.begin(),
                                lower_bound(table.begin(), table.end(), f + 1));
      freq[val]--;
      table[tpos]--;

      const int rank = index[tpos];
      st.update(rank, {f, rank * 1LL * f});

      const node left = st.query(0, rank);
      const node right = st.query(rank, kDistinct - 1);

      diversity -= (rank * left.val) - left.seq;
      diversity -= right.seq - (rank * right.val);

//      cout << "Diversity after removal : " << diversity << '\n';
    };

    for (int i = 0, l = 0, r = -1; i < queries.size(); i++) {
      while (r < queries[i].r) add(++r);
      while (r > queries[i].r) rem(r--);
      while (l < queries[i].l) rem(l++);
      while (l > queries[i].l) add(--l);
      const int x = queries[i].l;
      const int y = queries[i].r;
      const int64 m = (y - x + 1);
      queries[i].ans = diversity + (m * (m + 1) / 2);
    }

    sort(queries.begin(), queries.end(), [](const query &q1, const query &q2) {
      return q1.t < q2.t;
    });

    vector<int64> ans(queries.size());
    for (int i = 0; const query &q : queries) {
      ans[i++] = q.ans;
    }
    return ans;
  }
};

inline void solution() {
  const int n = Read::Int();
  const int q = Read::Int();
  vector<int> a = Read::Vec<int>(n);

  mo_s mo(a);
  for (int query = 0; query < q; query++) {
    const int l = Read::Int() - 1;
    const int r = Read::Int() - 1;
    mo.add_query(l, r, 0, query);
  }

  vector<int64> ans = mo.solve();
  for (const int64 x : ans) {
    cout << x << '\n';
  }
}

int32_t main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

  //freopen(".in", "r", stdin);
  //freopen(".out", "w", stdout);

  cout << fixed << setprecision(10);

  int testcases = 1;
  //cin >> testcases;
  while (testcases--) {
    solution();
  }
}

Compilation message (stderr)

diversity.cpp: In function 'std::vector<int> precalc(int)':
diversity.cpp:83:19: warning: range-based 'for' loops with initializer only available with '-std=c++2a' or '-std=gnu++2a'
   83 |   for (int i = 0; const int x : d) {
      |                   ^~~~~
diversity.cpp: In member function 'std::vector<long long int> mo_s::solve()':
diversity.cpp:315:38: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<mo_s::query>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  315 |     for (int i = 0, l = 0, r = -1; i < queries.size(); i++) {
      |                                    ~~^~~~~~~~~~~~~~~~
diversity.cpp:331:21: warning: range-based 'for' loops with initializer only available with '-std=c++2a' or '-std=gnu++2a'
  331 |     for (int i = 0; const query &q : queries) {
      |                     ^~~~~

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