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

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

using ll = long long;
#define int ll

#define rng(i, a, b) for (int i = int(a); i < int(b); i++)
#define rep(i, b) rng(i, 0, b)
#define gnr(i, a, b) for (int i = int(b) - 1; i >= int(a); i--)
#define per(i, b) gnr(i, 0, b)
#define pb push_back
#define eb emplace_back
#define a first
#define b second
#define bg begin()
#define ed end()
#define all(x) x.bg, x.ed
#define si(x) int(x.size())
#ifdef LOCAL
#define dmp(x) cerr << __LINE__ << " " << #x << " " << x << endl
#else
#define dmp(x) void(0)
#endif

template <class t, class u> void chmax(t &a, u b) {
  if (a < b)
    a = b;
}
template <class t, class u> void chmin(t &a, u b) {
  if (b < a)
    a = b;
}

template <class t> using vc = vector<t>;
template <class t> using vvc = vc<vc<t>>;

using pi = pair<int, int>;
using vi = vc<int>;

template <class t, class u>
ostream &operator<<(ostream &os, const pair<t, u> &p) {
  return os << "{" << p.a << "," << p.b << "}";
}

template <class t> ostream &operator<<(ostream &os, const vc<t> &v) {
  os << "{";
  for (auto e : v)
    os << e << ",";
  return os << "}";
}

#define mp make_pair
#define mt make_tuple
#define one(x) memset(x, -1, sizeof(x))
#define zero(x) memset(x, 0, sizeof(x))
#ifdef LOCAL
void dmpr(ostream &os) { os << endl; }
template <class T, class... Args>
void dmpr(ostream &os, const T &t, const Args &...args) {
  os << t << " ";
  dmpr(os, args...);
}
#define dmp2(...) dmpr(cerr, __LINE__, ##__VA_ARGS__)
#else
#define dmp2(...) void(0)
#endif

using uint = unsigned;
using ull = unsigned long long;

template <class t, size_t n>
ostream &operator<<(ostream &os, const array<t, n> &a) {
  return os << vc<t>(all(a));
}

template <int i, class T> void print_tuple(ostream &, const T &) {}

template <int i, class T, class H, class... Args>
void print_tuple(ostream &os, const T &t) {
  if (i)
    os << ",";
  os << get<i>(t);
  print_tuple<i + 1, T, Args...>(os, t);
}

template <class... Args>
ostream &operator<<(ostream &os, const tuple<Args...> &t) {
  os << "{";
  print_tuple<0, tuple<Args...>, Args...>(os, t);
  return os << "}";
}

template <class t> void print(t x, int suc = 1) {
  cout << x;
  if (suc == 1)
    cout << "\\n";
  if (suc == 2)
    cout << " ";
}

ll read() {
  ll i;
  cin >> i;
  return i;
}

vi readvi(int n, int off = 0) {
  vi v(n);
  rep(i, n) v[i] = read() + off;
  return v;
}

template <class T> void print(const vector<T> &v, int suc = 1) {
  rep(i, v.size()) print(v[i], i == int(v.size()) - 1 ? suc : 2);
}

string readString() {
  string s;
  cin >> s;
  return s;
}

template <class T> T sq(const T &t) { return t * t; }

// #define CAPITAL
void yes(bool ex = true) {
#ifdef CAPITAL
  cout << "YES"
       << "\\n";
#else
  cout << "Yes"
       << "\\n";
#endif
  if (ex)
    exit(0);
}
void no(bool ex = true) {
#ifdef CAPITAL
  cout << "NO"
       << "\\n";
#else
  cout << "No"
       << "\\n";
#endif
  if (ex)
    exit(0);
}
void possible(bool ex = true) {
#ifdef CAPITAL
  cout << "POSSIBLE"
       << "\\n";
#else
  cout << "Possible"
       << "\\n";
#endif
  if (ex)
    exit(0);
}
void impossible(bool ex = true) {
#ifdef CAPITAL
  cout << "IMPOSSIBLE"
       << "\\n";
#else
  cout << "Impossible"
       << "\\n";
#endif
  if (ex)
    exit(0);
}

constexpr ll ten(int n) { return n == 0 ? 1 : ten(n - 1) * 10; }

const ll infLL = LLONG_MAX / 3;

#ifdef int
const int inf = infLL;
#else
const int inf = INT_MAX / 2 - 100;
#endif

int topbit(signed t) { return t == 0 ? -1 : 31 - __builtin_clz(t); }
int topbit(ll t) { return t == 0 ? -1 : 63 - __builtin_clzll(t); }
int botbit(signed a) { return a == 0 ? 32 : __builtin_ctz(a); }
int botbit(ll a) { return a == 0 ? 64 : __builtin_ctzll(a); }
int popcount(signed t) { return __builtin_popcount(t); }
int popcount(ll t) { return __builtin_popcountll(t); }
bool ispow2(int i) { return i && (i & -i) == i; }
ll mask(int i) { return (ll(1) << i) - 1; }

bool inc(int a, int b, int c) { return a <= b && b <= c; }

template <class t> void mkuni(vc<t> &v) {
  sort(all(v));
  v.erase(unique(all(v)), v.ed);
}

ll rand_int(ll l, ll r) { //[l, r]
#ifdef LOCAL
  static mt19937_64 gen;
#else
  static mt19937_64 gen(chrono::steady_clock::now().time_since_epoch().count());
#endif
  return uniform_int_distribution<ll>(l, r)(gen);
}

template <class t> void myshuffle(vc<t> &a) {
  rep(i, si(a)) swap(a[i], a[rand_int(0, i)]);
}

template <class t> int lwb(const vc<t> &v, const t &a) {
  return lower_bound(all(v), a) - v.bg;
}

bool can(int k, vi a, vi pre) {
  vi u(k);
  for (auto v : pre)
    u[v] = 1;
  vc<pi> z;
  rep(i, k) if (!u[i]) z.eb(a[i], i);
  sort(all(z), greater<pi>());
  for (auto w : z) {
    if (--a[w.b] < 0)
      return false;
    pre.pb(w.b);
  }
  vi mn(k + 1), mx(k + 1);
  mn[0] = *min_element(all(a));
  rep(i, k) { mn[i + 1] = min(mn[i], --a[pre[i]]); }
  mx[k] = *max_element(all(a));
  per(i, k) { mx[i] = max(mx[i + 1], ++a[pre[i]]); }
  rep(i, k + 1) if (mn[i] * 2 >= mx[i]) return true;
  return false;
}

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

  int k;
  cin >> k;
  vi a = readvi(k);
  int len = accumulate(all(a), 0);

  vi ans, cur;
  int ok = 0;
  rep(_, len) {
    dmp2(ans, cur, ok);
    bool found = false;
    rep(nx, k) if (a[nx]) {
      int p = find(all(cur), nx) - cur.bg;
      vi tmp;
      int waf;
      if (p == si(cur)) {
        tmp = cur;
        waf = ok;
      } else {
        if (ok <= p)
          continue;
        tmp = vi(cur.bg + p + 1, cur.ed);
        waf = ok - (p + 1);
      }
      tmp.pb(nx);
      if (si(tmp) == k)
        waf = k;
      a[nx]--;
      bool relka = false;
      rep(j, waf + 1) {
        if (inc(1, j, waf - 1))
          continue;
        if (can(k, a, vi(tmp.bg + j, tmp.ed))) {
          ans.pb(nx);
          cur = tmp;
          ok = waf;
          found = true;
          relka = true;
          break;
        }
      }
      if (relka)
        break;
      a[nx]++;
    }
    if (!found) {
      assert(_ == 0);
      print(-1);
      return 0;
    }
  }
  assert(si(cur) == k);
  assert(si(ans) == len);
  for (auto &v : ans)
    v++;
  print(ans);
}


Error: TLE

Fixed submission:

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

using ll = long long;
#define int ll

#define rng(i, a, b) for (int i = int(a); i < int(b); i++)
#define rep(i, b) rng(i, 0, b)
#define gnr(i, a, b) for (int i = int(b) - 1; i >= int(a); i--)
#define per(i, b) gnr(i, 0, b)
#define pb push_back
#define eb emplace_back
#define a first
#define b second
#define bg begin()
#define ed end()
#define all(x) x.bg, x.ed
#define si(x) int(x.size())
#ifdef LOCAL
#define dmp(x) cerr << __LINE__ << " " << #x << " " << x << endl
#else
#define dmp(x) void(0)
#endif

template <class t, class u> void chmax(t &a, u b) {
  if (a < b)
    a = b;
}
template <class t, class u> void chmin(t &a, u b) {
  if (b < a)
    a = b;
}

template <class t> using vc = vector<t>;
template <class t> using vvc = vc<vc<t>>;

using pi = pair<int, int>;
using vi = vc<int>;

template <class t, class u>
ostream &operator<<(ostream &os, const pair<t, u> &p) {
  return os << "{" << p.a << "," << p.b << "}";
}

template <class t> ostream &operator<<(ostream &os, const vc<t> &v) {
  os << "{";
  for (auto e : v)
    os << e << ",";
  return os << "}";
}

#define mp make_pair
#define mt make_tuple
#define one(x) memset(x, -1, sizeof(x))
#define zero(x) memset(x, 0, sizeof(x))
#ifdef LOCAL
void dmpr(ostream &os) { os << endl; }
template <class T, class... Args>
void dmpr(ostream &os, const T &t, const Args &...args) {
  os << t << " ";
  dmpr(os, args...);
}
#define dmp2(...) dmpr(cerr, __LINE__, ##__VA_ARGS__)
#else
#define dmp2(...) void(0)
#endif

using uint = unsigned;
using ull = unsigned long long;

template <class t, size_t n>
ostream &operator<<(ostream &os, const array<t, n> &a) {
  return os << vc<t>(all(a));
}

template <int i, class T> void print_tuple(ostream &, const T &) {}

template <int i, class T, class H, class... Args>
void print_tuple(ostream &os, const T &t) {
  if (i)
    os << ",";
  os << get<i>(t);
  print_tuple<i + 1, T, Args...>(os, t);
}

template <class... Args>
ostream &operator<<(ostream &os, const tuple<Args...> &t) {
  os << "{";
  print_tuple<0, tuple<Args...>, Args...>(os, t);
  return os << "}";
}

template <class t> void print(t x, int suc = 1) {
  cout << x;
  if (suc == 1)
    cout << "\\n";
  if (suc == 2)
    cout << " ";
}

ll read() {
  ll i;
  cin >> i;
  return i;
}

vi readvi(int n, int off = 0) {
  vi v(n);
  rep(i, n) v[i] = read() + off;
  return v;
}

template <class T> void print(const vector<T> &v, int suc = 1) {
  rep(i, v.size()) print(v[i], i == int(v.size()) - 1 ? suc : 2);
}

string readString() {
  string s;
  cin >> s;
  return s;
}

template <class T> T sq(const T &t) { return t * t; }

// #define CAPITAL
void yes(bool ex = true) {
#ifdef CAPITAL
  cout << "YES"
       << "\\n";
#else
  cout << "Yes"
       << "\\n";
#endif
  if (ex)
    exit(0);
}
void no(bool ex = true) {
#ifdef CAPITAL
  cout << "NO"
       << "\\n";
#else
  cout << "No"
       << "\\n";
#endif
  if (ex)
    exit(0);
}
void possible(bool ex = true) {
#ifdef CAPITAL
  cout << "POSSIBLE"
       << "\\n";
#else
  cout << "Possible"
       << "\\n";
#endif
  if (ex)
    exit(0);
}
void impossible(bool ex = true) {
#ifdef CAPITAL
  cout << "IMPOSSIBLE"
       << "\\n";
#else
  cout << "Impossible"
       << "\\n";
#endif
  if (ex)
    exit(0);
}

constexpr ll ten(int n) { return n == 0 ? 1 : ten(n - 1) * 10; }

const ll infLL = LLONG_MAX / 3;

#ifdef int
const int inf = infLL;
#else
const int inf = INT_MAX / 2 - 100;
#endif

int topbit(signed t) { return t == 0 ? -1 : 31 - __builtin_clz(t); }
int topbit(ll t) { return t == 0 ? -1 : 63 - __builtin_clzll(t); }
int botbit(signed a) { return a == 0 ? 32 : __builtin_ctz(a); }
int botbit(ll a) { return a == 0 ? 64 : __builtin_ctzll(a); }
int popcount(signed t) { return __builtin_popcount(t); }
int popcount(ll t) { return __builtin_popcountll(t); }
bool ispow2(int i) { return i && (i & -i) == i; }
ll mask(int i) { return (ll(1) << i) - 1; }

bool inc(int a, int b, int c) { return a <= b && b <= c; }

template <class t> void mkuni(vc<t> &v) {
  sort(all(v));
  v.erase(unique(all(v)), v.ed);
}

ll rand_int(ll l, ll r) { //[l, r]
#ifdef LOCAL
  static mt19937_64 gen;
#else
  static mt19937_64 gen(chrono::steady_clock::now().time_since_epoch().count());
#endif
  return uniform_int_distribution<ll>(l, r)(gen);
}

template <class t> void myshuffle(vc<t> &a) {
  rep(i, si(a)) swap(a[i], a[rand_int(0, i)]);
}

template <class t> int lwb(const vc<t> &v, const t &a) {
  return lower_bound(all(v), a) - v.bg;
}

bool can(int k, vi a, vi pre) {
  vi u(k);
  for (auto v : pre)
    u[v] = 1;
  vc<pi> z;
  rep(i, k) if (!u[i]) z.eb(a[i], i);
  sort(all(z), greater<pi>());
  for (auto w : z) {
    if (--a[w.b] < 0)
      return false;
    pre.pb(w.b);
  }
  vi mn(k + 1), mx(k + 1);
  mn[0] = *min_element(all(a));
  rep(i, k) { mn[i + 1] = min(mn[i], --a[pre[i]]); }
  mx[k] = *max_element(all(a));
  per(i, k) { mx[i] = max(mx[i + 1], ++a[pre[i]]); }
  rep(i, k + 1) if (mn[i] * 2 >= mx[i]) return true;
  return false;
}

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

  int k;
  cin >> k;
  vi a = readvi(k);
  int len = accumulate(all(a), 0);

  vi ans, cur;
  int ok = 0;
  rep(_, len) {
    dmp2(ans, cur, ok);
    bool found = false;
    rep(nx, k) if (a[nx]) {
      int p = find(all(cur), nx) - cur.bg;
      vi tmp;
      int waf;
      if (p == si(cur)) {
        tmp = cur;
        waf = ok;
      } else {
        if (ok <= p)
          continue;
        tmp = vi(cur.bg + p + 1, cur.ed);
        waf = ok - (p + 1);
      }
      tmp.pb(nx);
      if (si(tmp) == k)
        waf = k;
      a[nx]--;
      bool relka = false;
      per(j, waf + 1) {
        if (can(k, a, vi(tmp.bg + j, tmp.ed))) {
          ans.pb(nx);
          cur = tmp;
          ok = waf;
          found = true;
          relka = true;
          break;
        }
      }
      if (relka)
        break;
      a[nx]++;
    }
    if (!found) {
      assert(_ == 0);
      print(-1);
      return 0;
    }
  }
  assert(si(cur) == k);
  assert(si(ans) == len);
  for (auto &v : ans)
    v++;
  print(ans);
}