Problem p03277 (C++)

Buggy submission (status: Time Limit Exceeded):
/**
 *
 * AtCoder Beginner Contest 107 - Problem D: Median of Medians
 * (https://beta.atcoder.jp/contests/abc107/tasks/arc107_b)
 *
 * Find the index j of the element i, if the vector is read backwards.
 *
 * Math solution: j = N - i + 1.
 */
#include <iostream>

#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstring>
#include <map>
#include <queue>
#include <set>
#include <vector>

using namespace std;

using ll = long long;
using ii = pair<ll, ll>;

#define REP(i, a, b) for (int(i) = (a); (i) <= (b); (i)++)

#define LOG(var, sep) (cerr << #var << " = " << (var) << (sep))
#define LOGM(msg) (cerr << (msg))
#define LOGV(vec, sep)                                                         \
  {                                                                            \
    cerr << #vec << " =";                                                      \
    for (const auto &v : (vec))                                                \
      cerr << " " << v;                                                        \
    cerr << (sep);                                                             \
  }
#define LOGA(arr, N, sep)                                                      \
  {                                                                            \
    cerr << #arr << " =";                                                      \
    for (int i = 0; i < N; ++i)                                                \
      cerr << " " << (arr)[i];                                                 \
    cerr << (sep);                                                             \
  }
#define LOGA2(arr, N, M, sep)                                                  \
  {                                                                            \
    cerr << #arr << " =\n";                                                    \
    for (int i = 0; i < N; ++i) {                                              \
      for (int j = 0; j < M; ++j)                                              \
        cerr << (arr)[i][j] << " ";                                            \
      cerr << endl;                                                            \
    }                                                                          \
    cerr << (sep);                                                             \
  }

ostream &operator<<(ostream &os, const ii &p) {
  os << "(" << p.first << ", " << p.second << ")";
  return os;
}

const int MAX{2000010};

int as[MAX], xs[MAX], S[MAX], temp[MAX];

ll inversions(int a, int b) {
  if (b - a == 1)
    return 0;

  int m = b + (a - b) / 2;

  auto sum = inversions(a, m) + inversions(m, b);

  int x = a, y = m, idx = 0;

  while (x < m and y < b) {
    if (S[x] <= S[y])
      temp[idx++] = S[x++];
    else {
      temp[idx++] = S[y++];
      sum += (x - a);
    }
  }

  while (x < m) {
    temp[idx++] = S[x++];
  }

  while (y < b) {
    temp[idx++] = S[y++];
    sum += (m - a);
  }

  for (int i = a; i < b; ++i)
    S[i] = temp[i - a];

  // printf("--- inversao [%d,%d] = %lld\n", a, b, sum);

  return sum;
}

ll pairs_with_median_larger_than(int x, int N) {
  S[0] = 0;

  REP(i, 1, N)
  S[i] = S[i - 1] + (as[i] >= x ? 1 : -1);

  // LOGA(as, N + 1, '\n');
  // LOG(x, ' '); LOGA(S, N + 1, '\n');

  auto ans = inversions(0, N + 1);
  // LOG(ans, '\n');

  return ans;
}

void medians(int N) {
  map<int, int> ms;

  REP(i, 1, N)
  REP(j, i, N) {
    vector<int> ns;

    REP(k, i, j)
    ns.push_back(as[k]);

    sort(ns.begin(), ns.end());
    auto pos = (ns.size() / 2);
    auto median = ns[pos];
    // LOG(i, ' '); LOG(j, ' '); LOGV(ns, ' '); LOG(pos, ' '); LOG(median,
    // '\n');
    ms[median]++;
  }

  //    for (auto p : ms)
  //        LOG(p, '\n');
}

void solve(ll N) {
  set<int> ns;

  REP(i, 1, N)
  ns.insert(as[i]);

  vector<int> ys(ns.begin(), ns.end());

  int a = 0, b = ys.size() - 1, ans = -1;
  auto pairs = N * (N + 1) / 2;

  //    LOG(pairs, '\n');

  while (a <= b) {
    auto m = a + (b - a) / 2;

    if (pairs_with_median_larger_than(ys[m], N) >= pairs / 2) {
      ans = max(ans, ys[m]);
      a = m + 1;
    } else
      b = m - 1;
  }

  cout << ans << endl;
}

int main() {
  ios::sync_with_stdio(false);

  int N;
  cin >> N;

  REP(i, 1, N)
  cin >> as[i];

  medians(N);
  solve(N);

  return 0;
}


Fixed submission:
/**
 *
 * AtCoder Beginner Contest 107 - Problem D: Median of Medians
 * (https://beta.atcoder.jp/contests/abc107/tasks/arc107_b)
 *
 * Find the index j of the element i, if the vector is read backwards.
 *
 * Math solution: j = N - i + 1.
 */
#include <iostream>

#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstring>
#include <map>
#include <queue>
#include <set>
#include <vector>

using namespace std;

using ll = long long;
using ii = pair<ll, ll>;

#define REP(i, a, b) for (int(i) = (a); (i) <= (b); (i)++)

#define LOG(var, sep) (cerr << #var << " = " << (var) << (sep))
#define LOGM(msg) (cerr << (msg))
#define LOGV(vec, sep)                                                         \
  {                                                                            \
    cerr << #vec << " =";                                                      \
    for (const auto &v : (vec))                                                \
      cerr << " " << v;                                                        \
    cerr << (sep);                                                             \
  }
#define LOGA(arr, N, sep)                                                      \
  {                                                                            \
    cerr << #arr << " =";                                                      \
    for (int i = 0; i < N; ++i)                                                \
      cerr << " " << (arr)[i];                                                 \
    cerr << (sep);                                                             \
  }
#define LOGA2(arr, N, M, sep)                                                  \
  {                                                                            \
    cerr << #arr << " =\n";                                                    \
    for (int i = 0; i < N; ++i) {                                              \
      for (int j = 0; j < M; ++j)                                              \
        cerr << (arr)[i][j] << " ";                                            \
      cerr << endl;                                                            \
    }                                                                          \
    cerr << (sep);                                                             \
  }

ostream &operator<<(ostream &os, const ii &p) {
  os << "(" << p.first << ", " << p.second << ")";
  return os;
}

const int MAX{2000010};

int as[MAX], xs[MAX], S[MAX], temp[MAX];

ll inversions(int a, int b) {
  if (b - a == 1)
    return 0;

  int m = b + (a - b) / 2;

  auto sum = inversions(a, m) + inversions(m, b);

  int x = a, y = m, idx = 0;

  while (x < m and y < b) {
    if (S[x] <= S[y])
      temp[idx++] = S[x++];
    else {
      temp[idx++] = S[y++];
      sum += (x - a);
    }
  }

  while (x < m) {
    temp[idx++] = S[x++];
  }

  while (y < b) {
    temp[idx++] = S[y++];
    sum += (m - a);
  }

  for (int i = a; i < b; ++i)
    S[i] = temp[i - a];

  // printf("--- inversao [%d,%d] = %lld\n", a, b, sum);

  return sum;
}

ll pairs_with_median_larger_than(int x, int N) {
  S[0] = 0;

  REP(i, 1, N)
  S[i] = S[i - 1] + (as[i] >= x ? 1 : -1);

  // LOGA(as, N + 1, '\n');
  // LOG(x, ' '); LOGA(S, N + 1, '\n');

  auto ans = inversions(0, N + 1);
  // LOG(ans, '\n');

  return ans;
}

void medians(int N) {
  map<int, int> ms;

  REP(i, 1, N)
  REP(j, i, N) {
    vector<int> ns;

    REP(k, i, j)
    ns.push_back(as[k]);

    sort(ns.begin(), ns.end());
    auto pos = (ns.size() / 2);
    auto median = ns[pos];
    // LOG(i, ' '); LOG(j, ' '); LOGV(ns, ' '); LOG(pos, ' '); LOG(median,
    // '\n');
    ms[median]++;
  }

  //    for (auto p : ms)
  //        LOG(p, '\n');
}

void solve(ll N) {
  set<int> ns;

  REP(i, 1, N)
  ns.insert(as[i]);

  vector<int> ys(ns.begin(), ns.end());

  int a = 0, b = ys.size() - 1, ans = -1;
  auto pairs = N * (N + 1) / 2;

  //    LOG(pairs, '\n');

  while (a <= b) {
    auto m = a + (b - a) / 2;

    if (pairs_with_median_larger_than(ys[m], N) >= pairs / 2) {
      ans = max(ans, ys[m]);
      a = m + 1;
    } else
      b = m - 1;
  }

  cout << ans << endl;
}

int main() {
  ios::sync_with_stdio(false);

  int N;
  cin >> N;

  REP(i, 1, N)
  cin >> as[i];

  //    medians(N);
  solve(N);

  return 0;
}