Problem p02878 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
#define mp make_pair
#define PI pair<ll, ll>
#define poly vector<ll>
#define mem(a) memset((a), 0, sizeof(a))
#define For(i, l, r) for (int i = (int)(l); i <= (int)(r); i++)
#define Rep(i, r, l) for (int i = (int)(r); i >= (int)(l); i--)
#define pb push_back
#define fi first
#define se second
#define SZ(x) ((int)(x.size()))
inline char gc() {
  static char buf[100000], *p1 = buf, *p2 = buf;
  return p1 == p2 && (p2 = (p1 = buf) + fread(buf, 1, 100000, stdin), p1 == p2)
             ? EOF
             : *p1++;
}
#define gc getchar
inline ll read() {
  ll x = 0;
  char ch = gc();
  bool positive = 1;
  for (; !isdigit(ch); ch = gc())
    if (ch == '-')
      positive = 0;
  for (; isdigit(ch); ch = gc())
    x = x * 10 + ch - '0';
  return positive ? x : -x;
}
inline void write(ll a) {
  if (a < 0) {
    a = -a;
    putchar('-');
  }
  if (a >= 10)
    write(a / 10);
  putchar('0' + a % 10);
}
inline void writeln(ll a) {
  write(a);
  puts("");
}
inline void wri(ll a) {
  write(a);
  putchar(' ');
}
ll rnd() {
  ull ans = 0;
  For(i, 0, 4) ans = ans << 15 ^ rand();
  return ans % ((ull)1 << 63);
}
const int mod = 998244353, N = 10000005;
ll fac[N], ni[N];
ll ksm(ll a, int b) {
  int ans = 1;
  for (; b; b >>= 1) {
    if (b & 1)
      ans = ans * a % mod;
    a = a * a % mod;
  }
  return ans;
}

ll c(int a, int b) {
  return a >= b && b >= 0 ? fac[a] * ni[b] % mod * ni[a - b] % mod : 0;
}
ll g(int a, int b) {
  if (!b)
    return a == 0;
  return c(a + b - 1, b - 1);
}
ll f(int a, int b) {
  assert((a + b) % 2 == 0);
  return c(a, (a + b) / 2);
}
int main() {
#ifdef Brollan
  freopen("1.in", "r", stdin);
#endif

  int n = read(), a = read(), b = read();
  if (!a) {
    puts("1");
    return 0;
  }
  For(i, fac[0] = 1, n) fac[i] = fac[i - 1] * i % mod;
  ni[n] = ksm(fac[n], mod - 2);
  Rep(i, n, 1) ni[i - 1] = ni[i] * i % mod;
  ll ans = 0;
  For(k, b, n) {
    int sn = b - (k - b);
    if (sn > 0 && k - b <= min(b - 1, a)) {
      ans = (ans + (k < n ? g(n - k - 1, sn - (b - a) + 1) : (sn == b - a)) *
                       (f(k - 1, sn - 1) - f(k - 1, sn + 1) + mod)) %
            mod;
      // cout<<sn<<" "<<n-k-1<<" "<<sn-(b-a)<<"
      // "<<(k<n?g(n-k-1,sn-(b-a)+1):(sn==b-a))<<" "<<f(k,sn)<<" "<<f(k,sn+2)<<"
      // "<<ans<<endl;
    }
  }
  cout << ans << endl;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
#define mp make_pair
#define PI pair<int, int>
#define poly vector<ll>
#define For(i, l, r) for (int i = (int)(l); i <= (int)(r); i++)
#define Rep(i, r, l) for (int i = (int)(r); i >= (int)(l); i--)
#define pb push_back
#define fi first
#define se second
inline char gc() {
  static char buf[100000], *p1 = buf, *p2 = buf;
  return p1 == p2 && (p2 = (p1 = buf) + fread(buf, 1, 100000, stdin), p1 == p2)
             ? EOF
             : *p1++;
}
#define gc getchar
inline ll read() {
  ll x = 0;
  char ch = gc();
  bool positive = 1;
  for (; !isdigit(ch); ch = gc())
    if (ch == '-')
      positive = 0;
  for (; isdigit(ch); ch = gc())
    x = x * 10 + ch - '0';
  return positive ? x : -x;
}
inline void write(ll a) {
  if (a < 0) {
    a = -a;
    putchar('-');
  }
  if (a >= 10)
    write(a / 10);
  putchar('0' + a % 10);
}
inline void writeln(ll a) {
  write(a);
  puts("");
}
inline void wri(ll a) {
  write(a);
  putchar(' ');
}
inline ull rnd() { return ((ull)rand() << 30 ^ rand()) << 4 | rand() % 4; }
const int N = 200005;
int k, A[N], lx[N], rx[N], a[N], nxt[N], pre[N];
ll ycl[N];
priority_queue<PI, vector<PI>, greater<PI>> q;
ll solve(int l, int r) {
  r = nxt[r];
  int dq = l;
  ll ans = 0;
  int sum = 0;
  For(i, 1, k - 1) {
    dq = nxt[dq];
    ans++;
    if (dq == r)
      return ans;
  }
  while (dq != r) {
    sum += lx[l];
    l = nxt[l];
    ans += (ll)sum * rx[dq] + 1;
    dq = nxt[dq];
  }
  return ans;
}
int main() {
  int n = read();
  k = read();
  For(i, k, n) ycl[i] = ycl[i - 1] + (i - k + 1);
  For(i, 1, n) {
    a[i] = read();
    lx[i] = 1;
    rx[i] = 1;
  }
  nxt[0] = 1;
  pre[n + 1] = n;
  For(i, 1, n) {
    pre[i] = i - 1;
    nxt[i] = i + 1;
    q.push(mp(a[i], i));
  }
  // For(i,1,n)wri(ycl[i]); puts("");
  ll ans = 0;
  while (q.size()) {
    PI t = q.top();
    q.pop();
    int l = t.se, r = l, cnt = 1;
    while (a[nxt[r]] == a[l]) {
      r = nxt[r];
      cnt++;
      q.pop();
    }
    // cout<<q.size()<<" "<<cnt<<" "<<t.fi<<" "<<t.se<<" "<<l<<" "<<r<<" "<<"
    // "<<ycl[cnt]<<endl;
    l = pre[l];
    r = nxt[r];
    ans += solve(nxt[l], pre[r]);
    if (cnt < k)
      continue;
    int zs = cnt, jb = l;
    zs -= k;
    For(i, 1, k) jb = nxt[jb];
    For(i, 1, cnt / k) {
      int meiju = min(k, zs + 1);
      rx[l + i] = 0;
      For(j, 1, meiju) {
        rx[l + i] += rx[jb];
        jb = nxt[jb];
      }
      zs -= k;
    }
    zs = cnt - k;
    jb = r;
    For(i, 1, k) jb = pre[jb];
    Rep(i, cnt / k, 1) {
      int meiju = min(k, zs + 1);
      A[i] = 0;
      For(j, 1, meiju) {
        A[i] += lx[jb];
        jb = pre[jb];
      }
      zs -= k;
    }
    For(i, 1, cnt / k) lx[l + i] = A[i];
    For(i, 1, cnt / k) {
      nxt[l + i] = l + i + 1;
      pre[l + i] = l + i - 1;
      a[l + i] = t.fi + 1;
      q.push(mp(a[l + i], l + i));
    }
    nxt[l] = l + 1;
    pre[r] = l + cnt / k;
    nxt[l + cnt / k] = r;
    // cout<<ans<<endl;
    ans -= solve(nxt[l], pre[r]);
    // For(i,1,cnt/k)wri(lx[l+i]); puts("");
    // cout<<ans<<" "<<nxt[jb]<<" "<<pre[r]<<endl;
  }
  cout << ans << endl;
}