Problem p03649 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;

// template {{{
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define mt make_tuple
#define lb lower_bound
#define ub upper_bound
#define f first
#define s second

#define sz(x) int((x).size())
#define all(x) (x).begin(), (x).end()

#define FOR(i, a, b) for (int i = (a); i < (b); i++)
#define F0R(i, a) for (int i = 0; i < (a); i++)
#define FORd(i, a, b) for (int i = (b)-1; i >= (a); i--)
#define F0Rd(i, a) for (int i = (a)-1; i >= 0; i--)
#define trav(a, x) for (auto &a : x)

using ll = long long;
using vi = vector<int>;
using vvi = vector<vi>;
using vll = vector<ll>;
using vvll = vector<vll>;
using vb = vector<bool>;
using vd = vector<double>;
using vs = vector<string>;

template <typename T> void ckmin(T &a, const T &b) { a = min(a, b); }
template <typename T> void ckmax(T &a, const T &b) { a = max(a, b); }

mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());

namespace __input {
template <class T1, class T2> void re(pair<T1, T2> &p);
template <class T> void re(vector<T> &a);
template <class T, size_t SZ> void re(array<T, SZ> &a);

template <class T> void re(T &x) { cin >> x; }
void re(double &x) {
  string t;
  re(t);
  x = stod(t);
}
template <class Arg, class... Args> void re(Arg &first, Args &...rest) {
  re(first);
  re(rest...);
}

template <class T1, class T2> void re(pair<T1, T2> &p) { re(p.f, p.s); }
template <class T> void re(vector<T> &a) { F0R(i, sz(a)) re(a[i]); }
template <class T, size_t SZ> void re(array<T, SZ> &a) { F0R(i, SZ) re(a[i]); }
} // namespace __input
using namespace __input;

namespace __output {
template <class T1, class T2> void pr(const pair<T1, T2> &x);
template <class T, size_t SZ> void pr(const array<T, SZ> &x);
template <class T> void pr(const vector<T> &x);
template <class T> void pr(const set<T> &x);
template <class T1, class T2> void pr(const map<T1, T2> &x);

template <class T> void pr(const T &x) { cout << x; }
template <class Arg, class... Args>
void pr(const Arg &first, const Args &...rest) {
  pr(first);
  pr(rest...);
}

template <class T1, class T2> void pr(const pair<T1, T2> &x) {
  pr("{", x.f, ", ", x.s, "}");
}
template <class T> void prContain(const T &x) {
  pr("{");
  bool fst = 1;
  for (const auto &a : x)
    pr(!fst ? ", " : "", a), fst = 0;
  pr("}");
}
template <class T, size_t SZ> void pr(const array<T, SZ> &x) { prContain(x); }
template <class T> void pr(const vector<T> &x) { prContain(x); }
template <class T> void pr(const set<T> &x) { prContain(x); }
template <class T1, class T2> void pr(const map<T1, T2> &x) { prContain(x); }

void ps() { pr("\n"); }
template <class Arg> void ps(const Arg &first) {
  pr(first);
  ps();
}
template <class Arg, class... Args>
void ps(const Arg &first, const Args &...rest) {
  pr(first, " ");
  ps(rest...);
}
} // namespace __output
using namespace __output;

#define TRACE(x) x
#define __pn(x) pr(#x, " = ")
#define pd(...) __pn((__VA_ARGS__)), ps(__VA_ARGS__)

namespace __io {
void setIn(string s) { freopen(s.c_str(), "r", stdin); }
void setOut(string s) { freopen(s.c_str(), "w", stdout); }
void setIO(string s = "") {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  if (sz(s)) {
    setIn(s + ".in"), setOut(s + ".out");
  }
}
} // namespace __io
using namespace __io;
// }}}

int main() {
  setIO();

  int N;
  re(N);
  vll a(N);
  re(a);
  sort(all(a), greater<ll>());

  ll ops = 0;
  while (a[0] >= 2 * N - 1) {
    int j = 0;
    while (j < N && a[j] > a[0] - N)
      j++;

    if (j == N) {
      ll r = a.back() - N + 1;
      ops += r * N;
      trav(v, a) v -= r;
      break;
    }

    ll gap = (a[0] - N + 1 - a[j]);
    ll moves = (gap + N + 2) / (N + 1);

    ops += moves * j;
    F0R(i, N) {
      if (i < j)
        a[i] -= moves * (N - j + 1);
      else
        a[i] += moves * j;
    }

    sort(all(a), greater<ll>());
  }

  while (a[0] >= N) {
    ops++;
    a[0] -= N;
    F0R(i, N - 1) a[i + 1]++;

    sort(all(a), greater<ll>());
  }

  ps(ops);

  return 0;
}


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

// template {{{
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define mt make_tuple
#define lb lower_bound
#define ub upper_bound
#define f first
#define s second

#define sz(x) int((x).size())
#define all(x) (x).begin(), (x).end()

#define FOR(i, a, b) for (int i = (a); i < (b); i++)
#define F0R(i, a) for (int i = 0; i < (a); i++)
#define FORd(i, a, b) for (int i = (b)-1; i >= (a); i--)
#define F0Rd(i, a) for (int i = (a)-1; i >= 0; i--)
#define trav(a, x) for (auto &a : x)

using ll = long long;
using vi = vector<int>;
using vvi = vector<vi>;
using vll = vector<ll>;
using vvll = vector<vll>;
using vb = vector<bool>;
using vd = vector<double>;
using vs = vector<string>;

template <typename T> void ckmin(T &a, const T &b) { a = min(a, b); }
template <typename T> void ckmax(T &a, const T &b) { a = max(a, b); }

mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());

namespace __input {
template <class T1, class T2> void re(pair<T1, T2> &p);
template <class T> void re(vector<T> &a);
template <class T, size_t SZ> void re(array<T, SZ> &a);

template <class T> void re(T &x) { cin >> x; }
void re(double &x) {
  string t;
  re(t);
  x = stod(t);
}
template <class Arg, class... Args> void re(Arg &first, Args &...rest) {
  re(first);
  re(rest...);
}

template <class T1, class T2> void re(pair<T1, T2> &p) { re(p.f, p.s); }
template <class T> void re(vector<T> &a) { F0R(i, sz(a)) re(a[i]); }
template <class T, size_t SZ> void re(array<T, SZ> &a) { F0R(i, SZ) re(a[i]); }
} // namespace __input
using namespace __input;

namespace __output {
template <class T1, class T2> void pr(const pair<T1, T2> &x);
template <class T, size_t SZ> void pr(const array<T, SZ> &x);
template <class T> void pr(const vector<T> &x);
template <class T> void pr(const set<T> &x);
template <class T1, class T2> void pr(const map<T1, T2> &x);

template <class T> void pr(const T &x) { cout << x; }
template <class Arg, class... Args>
void pr(const Arg &first, const Args &...rest) {
  pr(first);
  pr(rest...);
}

template <class T1, class T2> void pr(const pair<T1, T2> &x) {
  pr("{", x.f, ", ", x.s, "}");
}
template <class T> void prContain(const T &x) {
  pr("{");
  bool fst = 1;
  for (const auto &a : x)
    pr(!fst ? ", " : "", a), fst = 0;
  pr("}");
}
template <class T, size_t SZ> void pr(const array<T, SZ> &x) { prContain(x); }
template <class T> void pr(const vector<T> &x) { prContain(x); }
template <class T> void pr(const set<T> &x) { prContain(x); }
template <class T1, class T2> void pr(const map<T1, T2> &x) { prContain(x); }

void ps() { pr("\n"); }
template <class Arg> void ps(const Arg &first) {
  pr(first);
  ps();
}
template <class Arg, class... Args>
void ps(const Arg &first, const Args &...rest) {
  pr(first, " ");
  ps(rest...);
}
} // namespace __output
using namespace __output;

#define TRACE(x) x
#define __pn(x) pr(#x, " = ")
#define pd(...) __pn((__VA_ARGS__)), ps(__VA_ARGS__)

namespace __io {
void setIn(string s) { freopen(s.c_str(), "r", stdin); }
void setOut(string s) { freopen(s.c_str(), "w", stdout); }
void setIO(string s = "") {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  if (sz(s)) {
    setIn(s + ".in"), setOut(s + ".out");
  }
}
} // namespace __io
using namespace __io;
// }}}

int main() {
  setIO();

  int N;
  re(N);
  vll a(N);
  re(a);
  sort(all(a), greater<ll>());

  ll ops = 0;
  while (a[0] >= 2 * N - 1) {
    int j = 0;
    while (j < N && a[j] > a[0] - N)
      j++;

    if (j == N) {
      ll r = a.back() - N + 1;
      ops += r * N;
      trav(v, a) v -= r;
      break;
    }

    ll gap = (a[0] - N + 1 - a[j]);
    ll moves = (gap + N) / (N + 1);

    ops += moves * j;
    F0R(i, N) {
      if (i < j)
        a[i] -= moves * (N - j + 1);
      else
        a[i] += moves * j;
    }

    sort(all(a), greater<ll>());
  }

  while (a[0] >= N) {
    ops++;
    a[0] -= N;
    F0R(i, N - 1) a[i + 1]++;

    sort(all(a), greater<ll>());
  }

  ps(ops);

  return 0;
}