Problem p02652 (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 = 1e6 + 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 stor[1 << 20];

/**
 * Description: 1D range update and query, $SZ=2^p$.
 * Source: USACO Counting Haybales
 * Verification: SPOJ Horrible
 */

template <class T, int SZ> struct LazySeg {
  T mn[2 * SZ], mx[2 * SZ], lazy[2 * SZ];
  LazySeg() { F0R(i, 2 * SZ) mn[i] = mx[i] = lazy[i] = 0; }
  void push(int ind, int L, int R) { /// modify values for current node
    if (L != R)
      F0R(i, 2) lazy[2 * ind + i] += lazy[ind]; /// prop to children
    mn[ind] += lazy[ind];
    mx[ind] += lazy[ind];
    lazy[ind] = 0;
  } // recalc values for current node
  void pull(int ind) {
    mn[ind] = min(mn[2 * ind], mn[2 * ind + 1]);
    mx[ind] = max(mx[2 * ind], mx[2 * ind + 1]);
  }
  void build(int ind = 1, int L = 0, int R = SZ - 1) {
    lazy[ind] = 0;
    if (L == R) {
      mn[ind] = mx[ind] = stor[L];
      return;
    }
    int M = (L + R) / 2;
    build(2 * ind, L, M);
    build(2 * ind + 1, M + 1, R);
    pull(ind);
  }
  void upd(int lo, int hi, T inc, int ind = 1, int L = 0, int R = SZ - 1) {
    push(ind, L, R);
    if (hi < L || R < lo)
      return;
    if (lo <= L && R <= hi) {
      lazy[ind] = inc;
      push(ind, L, R);
      return;
    }
    int M = (L + R) / 2;
    upd(lo, hi, inc, 2 * ind, L, M);
    upd(lo, hi, inc, 2 * ind + 1, M + 1, R);
    pull(ind);
  }
  T qmax(int lo, int hi, int ind = 1, int L = 0, int R = SZ - 1) {
    push(ind, L, R);
    if (lo > R || L > hi)
      return -MOD;
    if (lo <= L && R <= hi)
      return mx[ind];
    int M = (L + R) / 2;
    return max(qmax(lo, hi, 2 * ind, L, M),
               qmax(lo, hi, 2 * ind + 1, M + 1, R));
  }
  T qmin(int lo, int hi, int ind = 1, int L = 0, int R = SZ - 1) {
    push(ind, L, R);
    if (lo > R || L > hi)
      return MOD;
    if (lo <= L && R <= hi)
      return mn[ind];
    int M = (L + R) / 2;
    return min(qmin(lo, hi, 2 * ind, L, M),
               qmin(lo, hi, 2 * ind + 1, M + 1, R));
  }
  int getMin(T x, int ind = 1, int L = 0, int R = SZ - 1) { // max <= x
    push(ind, L, R);
    if (mx[ind] <= x)
      return L; // OK
    if (L == R)
      return MOD;
    int M = (L + R) / 2;
    int t = getMin(x, 2 * ind + 1, M + 1, R);
    if (t > M + 1)
      return t;
    int u = getMin(x, 2 * ind, L, M);
    return u == MOD ? t : u;
  }
};

LazySeg<int, 1 << 20> L;

str S;

/*int getMin(int x) {
        int lo = 0, hi = sz(S)+1;
        while (lo < hi) {
                int mid = (lo+hi)/2;
                if (L.qmax(mid,sz(S)) > x) lo = mid+1;
                else hi = mid;
        }
        return lo;
}*/

int ans = MOD;

void tri(int mx) {
  L.build(1);
  set<int> change;
  F0R(i, sz(S)) if (S[i] == '?') change.insert(i);
  for (;; mx += 2) {
    while (sz(change)) {
      int ind = L.getMin(mx - 2) - 1;
      auto it = change.lb(ind);
      if (it == end(change))
        break;
      int i = *it;
      change.erase(it);
      L.upd(i + 1, (1 << 20) - 1, 2);
    }
    ckmin(ans, mx - L.qmin(0, (1 << 20) - 1));
    if (!sz(change))
      return;
  }
}

int main() {
  setIO();
  re(S);
  {
    int bal = 0;
    F0R(i, sz(S)) {
      if (S[i] == '0' || S[i] == '?')
        bal--;
      else
        bal++;
      stor[i + 1] = bal;
    }
    FOR(i, sz(S) + 1, 1 << 20) stor[i] = stor[sz(S)];
  }
  int bal = 0, mx = 0;
  trav(c, S) {
    if (c == '0' || c == '?')
      bal--;
    else
      bal++;
    ckmax(mx, bal);
  }
  tri(mx);
  tri(mx + 1);
  ps(ans);
  // 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 = 1e6 + 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
}

/**
 * Description: 1D range update and query, $SZ=2^p$.
 * Source: USACO Counting Haybales
 * Verification: SPOJ Horrible
 */

template <class T, int SZ> struct LazySeg {
  T mn[2 * SZ], mx[2 * SZ], lazy[2 * SZ];
  LazySeg() { F0R(i, 2 * SZ) mn[i] = mx[i] = lazy[i] = 0; }
  void push(int ind, int L, int R) { /// modify values for current node
    if (L != R)
      F0R(i, 2) lazy[2 * ind + i] += lazy[ind]; /// prop to children
    mn[ind] += lazy[ind];
    mx[ind] += lazy[ind];
    lazy[ind] = 0;
  } // recalc values for current node
  void pull(int ind) {
    mn[ind] = min(mn[2 * ind], mn[2 * ind + 1]);
    mx[ind] = max(mx[2 * ind], mx[2 * ind + 1]);
  }
  void build() { ROF(i, 1, SZ) pull(i); }
  void upd(int lo, int hi, T inc, int ind = 1, int L = 0, int R = SZ - 1) {
    push(ind, L, R);
    if (hi < L || R < lo)
      return;
    if (lo <= L && R <= hi) {
      lazy[ind] = inc;
      push(ind, L, R);
      return;
    }
    int M = (L + R) / 2;
    upd(lo, hi, inc, 2 * ind, L, M);
    upd(lo, hi, inc, 2 * ind + 1, M + 1, R);
    pull(ind);
  }
  T qmax(int lo, int hi, int ind = 1, int L = 0, int R = SZ - 1) {
    push(ind, L, R);
    if (lo > R || L > hi)
      return -MOD;
    if (lo <= L && R <= hi)
      return mx[ind];
    int M = (L + R) / 2;
    return max(qmax(lo, hi, 2 * ind, L, M),
               qmax(lo, hi, 2 * ind + 1, M + 1, R));
  }
  T qmin(int lo, int hi, int ind = 1, int L = 0, int R = SZ - 1) {
    push(ind, L, R);
    if (lo > R || L > hi)
      return MOD;
    if (lo <= L && R <= hi)
      return mn[ind];
    int M = (L + R) / 2;
    return min(qmin(lo, hi, 2 * ind, L, M),
               qmin(lo, hi, 2 * ind + 1, M + 1, R));
  }
  int getMin(T x, int ind = 1, int L = 0, int R = SZ - 1) { // max <= x
    push(ind, L, R);
    if (mx[ind] <= x)
      return L; // OK
    if (L == R)
      return MOD;
    int M = (L + R) / 2;
    int t = getMin(x, 2 * ind + 1, M + 1, R);
    if (t > M + 1)
      return t;
    int u = getMin(x, 2 * ind, L, M);
    return u == MOD ? t : u;
  }
};

LazySeg<int, 1 << 20> L;

str S;

/*int getMin(int x) {
        int lo = 0, hi = sz(S)+1;
        while (lo < hi) {
                int mid = (lo+hi)/2;
                if (L.qmax(mid,sz(S)) > x) lo = mid+1;
                else hi = mid;
        }
        return lo;
}*/

int ans = MOD;

void tri(int mx) {
  L = LazySeg<int, 1 << 20>();
  F0R(i, sz(S)) {
    if (S[i] == '0' || S[i] == '?')
      L.upd(i + 1, (1 << 20) - 1, -1);
    else
      L.upd(i + 1, (1 << 20) - 1, 1);
  }
  set<int> change;
  F0R(i, sz(S)) if (S[i] == '?') change.insert(i);
  for (;; mx += 2) {
    while (sz(change)) {
      int ind = L.getMin(mx - 2) - 1;
      auto it = change.lb(ind);
      if (it == end(change))
        break;
      int i = *it;
      change.erase(it);
      L.upd(i + 1, (1 << 20) - 1, 2);
    }
    ckmin(ans, mx - L.qmin(0, (1 << 20) - 1));
    if (!sz(change))
      return;
  }
}

int main() {
  setIO();
  re(S);
  int bal = 0, mx = 0;
  trav(c, S) {
    if (c == '0' || c == '?')
      bal--;
    else
      bal++;
    ckmax(mx, bal);
  }
  tri(mx);
  tri(mx + 1);
  ps(ans);
  // 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
 */