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

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

typedef long long ll;
typedef long double ld;
typedef double db;
typedef string str;

typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef pair<db, db> pd;

typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<db> vd;
typedef vector<str> vs;
typedef vector<pi> vpi;
typedef vector<pl> vpl;
typedef vector<pd> vpd;

#define mp make_pair
#define f first
#define s second
#define sz(x) (int)x.size()
#define all(x) begin(x), end(x)
#define rall(x) (x).rbegin(), (x).rend()
#define rsz resize
#define ins insert
#define ft front()
#define bk back()
#define pf push_front
#define pb push_back
#define eb emplace_back
#define lb lower_bound
#define ub upper_bound

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

const int MOD = 1e9 + 7; // 998244353;
const int MX = 2e5 + 5;
const ll INF = 1e18;
const ld PI = acos((ld)-1);
const int xd[4] = {1, 0, -1, 0}, yd[4] = {0, 1, 0, -1};
mt19937 rng((uint32_t)chrono::steady_clock::now().time_since_epoch().count());

template <class T> bool ckmin(T &a, const T &b) { return b < a ? a = b, 1 : 0; }
template <class T> bool ckmax(T &a, const T &b) { return a < b ? a = b, 1 : 0; }
int pct(int x) { return __builtin_popcount(x); }
int bits(int x) { return 31 - __builtin_clz(x); } // floor(log2(x))
int cdiv(int a, int b) {
  return a / b + !(a < 0 || a % b == 0);
} // division of a by b rounded up, assumes b > 0
int fstTrue(function<bool(int)> f, int lo, int hi) {
  hi++;
  assert(lo <= hi); // assuming f is increasing
  while (lo < hi) { // find first index such that f is true
    int mid = (lo + hi) / 2;
    f(mid) ? hi = mid : lo = mid + 1;
  }
  return lo;
}

// INPUT
template <class A> void re(complex<A> &c);
template <class A, class B> void re(pair<A, B> &p);
template <class A> void re(vector<A> &v);
template <class A, size_t SZ> void re(array<A, SZ> &a);

template <class T> void re(T &x) { cin >> x; }
void re(db &d) {
  str t;
  re(t);
  d = stod(t);
}
void re(ld &d) {
  str t;
  re(t);
  d = stold(t);
}
template <class H, class... T> void re(H &h, T &...t) {
  re(h);
  re(t...);
}

template <class A> void re(complex<A> &c) {
  A a, b;
  re(a, b);
  c = {a, b};
}
template <class A, class B> void re(pair<A, B> &p) { re(p.f, p.s); }
template <class A> void re(vector<A> &x) { trav(a, x) re(a); }
template <class A, size_t SZ> void re(array<A, SZ> &x) { trav(a, x) re(a); }

// TO_STRING
#define ts to_string
str ts(char c) { return str(1, c); }
str ts(bool b) { return b ? "true" : "false"; }
str ts(const char *s) { return (str)s; }
str ts(str s) { return s; }
template <class A> str ts(complex<A> c) {
  stringstream ss;
  ss << c;
  return ss.str();
}
str ts(vector<bool> v) {
  str res = "{";
  F0R(i, sz(v)) res += char('0' + v[i]);
  res += "}";
  return res;
}
template <size_t SZ> str ts(bitset<SZ> b) {
  str res = "";
  F0R(i, SZ) res += char('0' + b[i]);
  return res;
}
template <class A, class B> str ts(pair<A, B> p);
template <class T> str ts(T v) { // containers with begin(), end()
  bool fst = 1;
  str res = "{";
  for (const auto &x : v) {
    if (!fst)
      res += ", ";
    fst = 0;
    res += ts(x);
  }
  res += "}";
  return res;
}
template <class A, class B> str ts(pair<A, B> p) {
  return "(" + ts(p.f) + ", " + ts(p.s) + ")";
}

// OUTPUT
template <class A> void pr(A x) { cout << ts(x); }
template <class H, class... T> void pr(const H &h, const T &...t) {
  pr(h);
  pr(t...);
}
void ps() { pr("\\n"); } // print w/ spaces
template <class H, class... T> void ps(const H &h, const T &...t) {
  pr(h);
  if (sizeof...(t))
    pr(" ");
  ps(t...);
}

// DEBUG
void DBG() { cerr << "]" << endl; }
template <class H, class... T> void DBG(H h, T... t) {
  cerr << ts(h);
  if (sizeof...(t))
    cerr << ", ";
  DBG(t...);
}
#ifdef LOCAL // compile with -DLOCAL
#define dbg(...)                                                               \\
  cerr << "LINE(" << __LINE__ << ") -> [" << #__VA_ARGS__ << "]: [",           \\
      DBG(__VA_ARGS__)
#else
#define dbg(...) 0
#endif

// FILE I/O
void setIn(string s) { freopen(s.c_str(), "r", stdin); }
void setOut(string s) { freopen(s.c_str(), "w", stdout); }
void unsyncIO() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
}
void setIO(string s = "") {
  unsyncIO();
  // cin.exceptions(cin.failbit);
  // throws exception when do smth illegal
  // ex. try to read letter into int
  if (sz(s)) {
    setIn(s + ".in"), setOut(s + ".out");
  } // for USACO
}

int K, sum;
vi _A, sol;

int getSum(vi v) {
  int res = 0;
  trav(t, v) res += t;
  return res;
}

bool initial(vi v) {
  vi A = _A;
  vector<bool> con(K);
  // dbg("OOPS",sz(A),K);
  trav(t, v) A[t]++, con[t] = 1;
  int lef = getSum(A);
  assert(lef >= K);
  int l = MOD, r = -MOD;
  trav(t, A) ckmin(l, t), ckmax(r, t);
  int lo = cdiv(lef, K), hi = lef / K * 2;
  dbg("UHOH", lef, v);
  FOR(num, lo, hi + 1) {
    int LO = (num + 1) / 2, HI = num;
    if (LO <= l && r <= HI) {
      if (num == 1 || LO < l || r < HI)
        return 1;
      bool bad = 0, flag = 0;
      trav(t, v) {
        flag |= A[t] == LO;
        if (A[t] == HI && flag) {
          bad = 1;
          break;
        }
      }
      if (flag) {
        F0R(i, K) if (!con[i] && A[i] == HI) {
          bad = 1;
          break;
        }
      }
      if (!bad) {
        dbg("OK", num, LO, HI);
        return 1;
      }
    }
  }
  return 0;
  /*vi con(K);
  trav(t,v) {
          assert(!con[t]);
          con[t] = 1;
  }
  F0R(i,K) if (!con[i]) {
          if (!A[i]) return 0;
          v.pb(i); A[i] --;
  }
  FOR(nex,1,K+1) {

  }
  dbg("AH",v,A); exit(0);*/
}

bool ok() {
  vi cnt(K);
  int ok = 0;
  vi cool(sz(sol) + 1);
  auto ad = [&](int r, int inc) {
    r = sol[r];
    ok -= cnt[r] == 1;
    cnt[r] += inc;
    ok += cnt[r] == 1;
  };
  int r = -1;
  F0R(l, sz(sol) - K + 1) {
    while (r < l + K - 1)
      ad(++r, 1);
    if (ok == K)
      cool[l]++, cool[r + 1]--;
    ad(l, -1);
  }
  FOR(i, 1, sz(sol) + 1) cool[i] += cool[i - 1];
  int fst = 0;
  while (cool[fst])
    fst++;
  if (sz(sol) == sum)
    return fst == sum;
  cnt = vi(K);
  FOR(i, fst, sz(sol)) if (cnt[sol[i]]++) return 0;
  dbg("HA", fst, sol);
  // fst to sz(sol)-1
  for (int i = fst; i + K >= sz(sol); --i) {
    // dbg("TRYING",i);
    if (i + K <= sum && initial(vi(begin(sol) + i, end(sol))))
      return 1;
    if (i == 0)
      break;
    if (cnt[sol[i - 1]]++)
      break;
  }
  return 0;
}

int main() {
  setIO();
  re(K);
  _A.rsz(K);
  re(_A);
  trav(t, _A) sum += t;
  if (!ok()) {
    ps(-1);
    exit(0);
  }
  while (sz(sol) < sum) {
    dbg("SOL", sol);
    bool done = 0;
    F0R(i, K) if (_A[i]) { // K*sum
      _A[i]--;
      sol.pb(i);
      if (ok()) {
        done = 1;
        break;
      }
      _A[i]++;
      sol.pop_back();
    }
    assert(done);
  }
  trav(t, sol) pr(t + 1, ' ');
  ps();
  // you should actually read the stuff at the bottom
}

/* stuff you should look for
 * int overflow, array bounds
 * special cases (n=1?)
 * do smth instead of nothing and stay organized
 * WRITE STUFF DOWN
 */


Error: TLE

Fixed submission:

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

typedef long long ll;
typedef long double ld;
typedef double db;
typedef string str;

typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef pair<db, db> pd;

typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<db> vd;
typedef vector<str> vs;
typedef vector<pi> vpi;
typedef vector<pl> vpl;
typedef vector<pd> vpd;

#define mp make_pair
#define f first
#define s second
#define sz(x) (int)x.size()
#define all(x) begin(x), end(x)
#define rall(x) (x).rbegin(), (x).rend()
#define rsz resize
#define ins insert
#define ft front()
#define bk back()
#define pf push_front
#define pb push_back
#define eb emplace_back
#define lb lower_bound
#define ub upper_bound

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

const int MOD = 1e9 + 7; // 998244353;
const int MX = 2e5 + 5;
const ll INF = 1e18;
const ld PI = acos((ld)-1);
const int xd[4] = {1, 0, -1, 0}, yd[4] = {0, 1, 0, -1};
mt19937 rng((uint32_t)chrono::steady_clock::now().time_since_epoch().count());

template <class T> bool ckmin(T &a, const T &b) { return b < a ? a = b, 1 : 0; }
template <class T> bool ckmax(T &a, const T &b) { return a < b ? a = b, 1 : 0; }
int pct(int x) { return __builtin_popcount(x); }
int bits(int x) { return 31 - __builtin_clz(x); } // floor(log2(x))
int cdiv(int a, int b) {
  return a / b + !(a < 0 || a % b == 0);
} // division of a by b rounded up, assumes b > 0
int fstTrue(function<bool(int)> f, int lo, int hi) {
  hi++;
  assert(lo <= hi); // assuming f is increasing
  while (lo < hi) { // find first index such that f is true
    int mid = (lo + hi) / 2;
    f(mid) ? hi = mid : lo = mid + 1;
  }
  return lo;
}

// INPUT
template <class A> void re(complex<A> &c);
template <class A, class B> void re(pair<A, B> &p);
template <class A> void re(vector<A> &v);
template <class A, size_t SZ> void re(array<A, SZ> &a);

template <class T> void re(T &x) { cin >> x; }
void re(db &d) {
  str t;
  re(t);
  d = stod(t);
}
void re(ld &d) {
  str t;
  re(t);
  d = stold(t);
}
template <class H, class... T> void re(H &h, T &...t) {
  re(h);
  re(t...);
}

template <class A> void re(complex<A> &c) {
  A a, b;
  re(a, b);
  c = {a, b};
}
template <class A, class B> void re(pair<A, B> &p) { re(p.f, p.s); }
template <class A> void re(vector<A> &x) { trav(a, x) re(a); }
template <class A, size_t SZ> void re(array<A, SZ> &x) { trav(a, x) re(a); }

// TO_STRING
#define ts to_string
str ts(char c) { return str(1, c); }
str ts(bool b) { return b ? "true" : "false"; }
str ts(const char *s) { return (str)s; }
str ts(str s) { return s; }
template <class A> str ts(complex<A> c) {
  stringstream ss;
  ss << c;
  return ss.str();
}
str ts(vector<bool> v) {
  str res = "{";
  F0R(i, sz(v)) res += char('0' + v[i]);
  res += "}";
  return res;
}
template <size_t SZ> str ts(bitset<SZ> b) {
  str res = "";
  F0R(i, SZ) res += char('0' + b[i]);
  return res;
}
template <class A, class B> str ts(pair<A, B> p);
template <class T> str ts(T v) { // containers with begin(), end()
  bool fst = 1;
  str res = "{";
  for (const auto &x : v) {
    if (!fst)
      res += ", ";
    fst = 0;
    res += ts(x);
  }
  res += "}";
  return res;
}
template <class A, class B> str ts(pair<A, B> p) {
  return "(" + ts(p.f) + ", " + ts(p.s) + ")";
}

// OUTPUT
template <class A> void pr(A x) { cout << ts(x); }
template <class H, class... T> void pr(const H &h, const T &...t) {
  pr(h);
  pr(t...);
}
void ps() { pr("\\n"); } // print w/ spaces
template <class H, class... T> void ps(const H &h, const T &...t) {
  pr(h);
  if (sizeof...(t))
    pr(" ");
  ps(t...);
}

// DEBUG
void DBG() { cerr << "]" << endl; }
template <class H, class... T> void DBG(H h, T... t) {
  cerr << ts(h);
  if (sizeof...(t))
    cerr << ", ";
  DBG(t...);
}
#ifdef LOCAL // compile with -DLOCAL
#define dbg(...)                                                               \\
  cerr << "LINE(" << __LINE__ << ") -> [" << #__VA_ARGS__ << "]: [",           \\
      DBG(__VA_ARGS__)
#else
#define dbg(...) 0
#endif

// FILE I/O
void setIn(string s) { freopen(s.c_str(), "r", stdin); }
void setOut(string s) { freopen(s.c_str(), "w", stdout); }
void unsyncIO() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
}
void setIO(string s = "") {
  unsyncIO();
  // cin.exceptions(cin.failbit);
  // throws exception when do smth illegal
  // ex. try to read letter into int
  if (sz(s)) {
    setIn(s + ".in"), setOut(s + ".out");
  } // for USACO
}

int K, sum;
vi _A, sol;

int getSum(vi v) {
  int res = 0;
  trav(t, v) res += t;
  return res;
}

bool initial(vi v) {
  vi A = _A;
  vector<bool> con(K);
  // dbg("OOPS",sz(A),K);
  trav(t, v) A[t]++, con[t] = 1;
  vpi oops;
  F0R(i, K) if (!con[i]) oops.pb({A[i], i});
  sort(rall(oops));
  trav(t, oops) v.pb(t.s);
  assert(sz(v) == K);
  int lef = getSum(A);
  assert(lef >= K);
  int l = MOD, r = -MOD;
  trav(t, A) ckmin(l, t), ckmax(r, t);
  int lo = cdiv(lef, K), hi = lef / K * 2;
  dbg("UHOH", lef, v);
  FOR(num, lo, hi + 1) {
    int LO = (num + 1) / 2, HI = num;
    if (LO <= l && r <= HI) {
      bool bad = 0;
      if (num > 1) {
        bool flag = 0;
        trav(t, v) {
          flag |= A[t] == LO;
          if (A[t] == HI && flag)
            bad = 1;
        }
      }
      if (!bad) {
        dbg("OK", num, LO, HI);
        return 1;
      }
    }
  }
  return 0;
  /*vi con(K);
  trav(t,v) {
          assert(!con[t]);
          con[t] = 1;
  }
  F0R(i,K) if (!con[i]) {
          if (!A[i]) return 0;
          v.pb(i); A[i] --;
  }
  FOR(nex,1,K+1) {

  }
  dbg("AH",v,A); exit(0);*/
}

bool ok() {
  vi cnt(K);
  int ok = 0;
  vi cool(sz(sol) + 1);
  auto ad = [&](int r, int inc) {
    r = sol[r];
    ok -= cnt[r] == 1;
    cnt[r] += inc;
    ok += cnt[r] == 1;
  };
  int r = -1;
  F0R(l, sz(sol) - K + 1) {
    while (r < l + K - 1)
      ad(++r, 1);
    if (ok == K)
      cool[l]++, cool[r + 1]--;
    ad(l, -1);
  }
  FOR(i, 1, sz(sol) + 1) cool[i] += cool[i - 1];
  int fst = 0;
  while (cool[fst])
    fst++;
  if (sz(sol) == sum)
    return fst == sum;
  cnt = vi(K);
  FOR(i, fst, sz(sol)) if (cnt[sol[i]]++) return 0;
  dbg("HA", fst, sol);
  // fst to sz(sol)-1
  for (int i = fst; i + K >= sz(sol); --i) {
    // dbg("TRYING",i);
    if (i + K <= sum && initial(vi(begin(sol) + i, end(sol))))
      return 1;
    if (i == 0)
      break;
    if (cnt[sol[i - 1]]++)
      break;
  }
  return 0;
}

int main() {
  setIO();
  re(K);
  _A.rsz(K);
  re(_A);
  trav(t, _A) sum += t;
  if (!ok()) {
    ps(-1);
    exit(0);
  }
  while (sz(sol) < sum) {
    dbg("SOL", sol);
    bool done = 0;
    F0R(i, K) if (_A[i]) {
      _A[i]--;
      sol.pb(i);
      if (ok()) {
        done = 1;
        break;
      }
      _A[i]++;
      sol.pop_back();
    }
    assert(done);
  }
  trav(t, sol) pr(t + 1, ' ');
  ps();
  // you should actually read the stuff at the bottom
}

/* stuff you should look for
 * int overflow, array bounds
 * special cases (n=1?)
 * do smth instead of nothing and stay organized
 * WRITE STUFF DOWN
 */