Problem p02615 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using Graph = vector<vector<int>>;
using WGraph = vector<vector<pair<int, ll>>>;

template <class T> inline bool chmax(T &a, const T &b) {
  if (b > a) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

constexpr int dx[4] = {-1, 0, 0, 1}, dy[4] = {0, -1, 1, 0};
constexpr int INF = 1e9;
constexpr int MOD = 1e9 + 7;
constexpr long long LINF = 1e18;
constexpr double EPS = 1e-10;
constexpr double PI = M_PI;

template <class RandomAccessIterator>
auto medianOfMedians(RandomAccessIterator first, RandomAccessIterator last,
                     int k) {
  auto _first = *first;
  using T = decltype(_first);

  int len = last - first;
  if (len == 1 && k == 0)
    return *first;

  int ms = (len + 4) / 5;
  T *medians = new T[ms];
  for (int i = 0; i < ms; ++i) {
    auto iter = first + 5 * i;
    if (5 * i + 4 < len) {
      for (int j = 0; j < 3; ++j) {
        int minInd = j;
        for (int k = j + 1; k < 5; ++k) {
          if (*(iter + k) < *(iter + minInd))
            minInd = k;
        }
        swap(*(iter + j), *(iter + minInd));
      }
      medians[i] = *(iter + 2);
    } else {
      medians[i] = *iter;
    }
  }

  T pivot = medianOfMedians(medians, medians + ms, ms / 2);
  delete[] medians;

  int store = 0, same = 0;
  for (int i = 0; i < len; ++i) {
    if (*(first + i) < pivot) {
      if (i != store + same)
        swap(*(first + store), *(first + store + same));
      swap(*(first + i), *(first + store));
      ++store;
    } else if (*(first + i) == pivot) {
      swap(*(first + i), *(first + store + same));
      ++same;
    }
  }

  if (store <= k && k < store + same) {
    return pivot;
  } else if (k < store) {
    return medianOfMedians(first, first + store, k);
  } else {
    return medianOfMedians(first + store + same, last, k - store - same);
  }
}

void solve() {
  int N;
  cin >> N;
  vector<int> a(N);
  for (int i = 0; i < N; ++i)
    cin >> a[i];

  int maxi = -INF;
  for (int x : a)
    chmax(maxi, x);
  int medi = medianOfMedians(a.begin(), a.end(), N / 2);
  ll res = 0;
  int cnt = -1;
  for (int x : a) {
    if (medi < x) {
      res += 2 * x;
      cnt += 2;
    }
  }
  res -= maxi;
  res += (ll)medi * (N - 1 - cnt);
  cout << res << '\\n';
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout << fixed << setprecision(20);

  solve();

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using Graph = vector<vector<int>>;
using WGraph = vector<vector<pair<int, ll>>>;

template <class T> inline bool chmax(T &a, const T &b) {
  if (b > a) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

constexpr int dx[4] = {-1, 0, 0, 1}, dy[4] = {0, -1, 1, 0};
constexpr int INF = 1e9;
constexpr int MOD = 1e9 + 7;
constexpr long long LINF = 1e18;
constexpr double EPS = 1e-10;
constexpr double PI = M_PI;

template <class RandomAccessIterator>
auto medianOfMedians(RandomAccessIterator first, RandomAccessIterator last,
                     int k) {
  auto _first = *first;
  using T = decltype(_first);

  int len = last - first;
  if (len == 1 && k == 0)
    return *first;

  int ms = (len + 4) / 5;
  T *medians = new T[ms];
  for (int i = 0; i < ms; ++i) {
    auto iter = first + 5 * i;
    if (5 * i + 4 < len) {
      for (int j = 0; j < 3; ++j) {
        int minInd = j;
        for (int k = j + 1; k < 5; ++k) {
          if (*(iter + k) < *(iter + minInd))
            minInd = k;
        }
        swap(*(iter + j), *(iter + minInd));
      }
      medians[i] = *(iter + 2);
    } else {
      medians[i] = *iter;
    }
  }

  T pivot = medianOfMedians(medians, medians + ms, ms / 2);
  delete[] medians;

  for (int i = 0; i < len; ++i) {
    if (*(first + i) == pivot) {
      swap(*(first + i), *(first + len - 1));
      break;
    }
  }

  int store = 0;
  for (int i = 0; i < len - 1; ++i) {
    if (*(first + i) < pivot) {
      swap(*(first + i), *(first + store));
      ++store;
    }
  }
  swap(*(first + store), *(first + len - 1));

  if (k == store) {
    return pivot;
  } else if (k < store) {
    return medianOfMedians(first, first + store, k);
  } else {
    return medianOfMedians(first + store + 1, last, k - store - 1);
  }
}

void solve() {
  int N;
  cin >> N;
  vector<int> a(N);
  for (int i = 0; i < N; ++i)
    cin >> a[i];

  int maxi = -INF;
  for (int x : a)
    chmax(maxi, x);
  int medi = medianOfMedians(a.begin(), a.end(), N / 2);
  ll res = 0;
  int cnt = -1;
  for (int x : a) {
    if (medi < x) {
      res += 2 * x;
      cnt += 2;
    }
  }
  res -= maxi;
  res += (ll)medi * (N - 1 - cnt);
  cout << res << '\\n';
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout << fixed << setprecision(20);

  solve();

  return 0;
}