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

// By TheOneYouWant
#pragma GCC optimize("-O2")
#include <atcoder/all>
#include <bits/stdc++.h>
using namespace std;
using namespace atcoder;
#define fastio                                                                 \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0)
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define all(x) x.begin(), x.end()
#define int long long int
#define memreset(a) memset(a, 0, sizeof(a))
#define testcase(t)                                                            \\
  int t;                                                                       \\
  cin >> t;                                                                    \\
  while (t--)
#define forstl(i, v) for (auto &i : v)
#define forn(i, e) for (int i = 0; i < e; ++i)
#define forsn(i, s, e) for (int i = s; i < e; ++i)
#define rforn(i, s) for (int i = s; i >= 0; --i)
#define rforsn(i, s, e) for (int i = s; i >= e; --i)
#define bitcount(a) __builtin_popcount(a) // set bits (add ll)
#define ln '\\n'
#define getcurrtime()                                                          \\
  cerr << "Time = " << ((double)clock() / CLOCKS_PER_SEC) << endl
#define dbgarr(v, s, e)                                                        \\
  cerr << #v << " = ";                                                         \\
  forsn(i, s, e) cerr << v[i] << ", ";                                         \\
  cerr << endl
#define inputfile freopen("input.txt", "r", stdin)
#define outputfile freopen("output.txt", "w", stdout)
#define dbg(args...)                                                           \\
  {                                                                            \\
    string _s = #args;                                                         \\
    replace(_s.begin(), _s.end(), ',', ' ');                                   \\
    stringstream _ss(_s);                                                      \\
    istream_iterator<string> _it(_ss);                                         \\
    err(_it, args);                                                            \\
  }
void err(istream_iterator<string> it) { cerr << endl; }
template <typename T, typename... Args>
void err(istream_iterator<string> it, T a, Args... args) {
  cerr << *it << " = " << a << "\\t";
  err(++it, args...);
}
template <typename T1, typename T2>
ostream &operator<<(ostream &c, pair<T1, T2> &v) {
  c << "(" << v.fi << "," << v.se << ")";
  return c;
}
template <template <class...> class TT, class... T>
ostream &operator<<(ostream &out, TT<T...> &c) {
  out << "{ ";
  forstl(x, c) out << x << " ";
  out << "}";
  return out;
}
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef pair<ll, ll> p64;
typedef pair<int, int> p32;
typedef pair<int, p32> p96;
typedef vector<ll> v64;
typedef vector<int> v32;
typedef vector<v32> vv32;
typedef vector<v64> vv64;
typedef vector<p32> vp32;
typedef vector<p64> vp64;
typedef vector<vp32> vvp32;
typedef map<int, int> m32;
const int LIM = 2e5 + 5, MOD = 998244353;
const ld EPS = 1e-9;

int read() {
  int xx = 0, ff = 1;
  char ch = getchar();
  while (ch > '9' || ch < '0') {
    if (ch == '-')
      ff = -1;
    ch = getchar();
  }
  while (ch >= '0' && ch <= '9') {
    xx = (xx << 3) + (xx << 1) + ch - '0';
    ch = getchar();
  }
  return xx * ff;
}
using mint = modint998244353;

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

struct S {
  mint a;
  int size;
};

struct F {
  mint a, b;
};

S op(S l, S r) { return S{l.a + r.a, l.size + r.size}; }

S e() { return S{0, 0}; }

S mapping(F l, S r) { return S{r.a * l.a + r.size * l.b, r.size}; }

F composition(F l, F r) { return F{r.a * l.a, r.b * l.a + l.b}; }

F id() { return F{1, 0}; }

ll mul(ll a, ll b) { return (a * b) % MOD; }
ll add(ll a, ll b) { return (a + b) % MOD; }

ll fastpow(ll a, ll p) {
  if (p == 0)
    return 1;
  ll z = fastpow(a, p / 2);
  z = mul(z, z);
  if (p % 2)
    z = mul(a, z);
  return z;
}

signed main() {

  int n, k;
  cin >> n >> k;
  int p[n];
  forn(i, n) {
    cin >> p[i];
    p[i]--;
  }

  vector<S> a(n);
  for (int i = 0; i < n; i++) {
    a[i] = S{0, 0};
  }

  lazy_segtree<S, op, e, F, mapping, composition, id> seg(a);

  mint ans = 0;

  for (int i = 0; i < k; i++) {
    mint add1 = seg.prod(0, p[i]).a;
    ans += add1;
    mint add2 = seg.prod(p[i] + 1, n).a;
    mint sz = seg.prod(p[i] + 1, n).size;
    add2 = sz - add2;
    ans += add2;
    seg.set(p[i], S{fastpow(2, MOD - 2), 1});
  }

  mint app = mul(k - 1, fastpow(k, MOD - 2));

  for (int i = k; i < n; i++) {
    seg.apply(0, n, F{app, 0});
    mint add1 = seg.prod(0, p[i]).a;
    ans += add1;
    mint add2 = seg.prod(p[i] + 1, n).a;
    mint sz = seg.prod(p[i] + 1, n).size;
    add2 = sz - add2;
    ans += add2;
    seg.set(p[i], S{fastpow(2, MOD - 2), 1});
  }
  cout << ans.val() << ln;
}

Error: TLE

Fixed submission:

// By TheOneYouWant
#pragma GCC optimize("-O2")
#include <bits/stdc++.h>
using namespace std;
#define fastio                                                                 \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0)
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define all(x) x.begin(), x.end()
#define memreset(a) memset(a, 0, sizeof(a))
#define testcase(t)                                                            \\
  int t;                                                                       \\
  cin >> t;                                                                    \\
  while (t--)
#define forstl(i, v) for (auto &i : v)
#define forn(i, e) for (int i = 0; i < e; ++i)
#define forsn(i, s, e) for (int i = s; i < e; ++i)
#define rforn(i, s) for (int i = s; i >= 0; --i)
#define rforsn(i, s, e) for (int i = s; i >= e; --i)
#define bitcount(a) __builtin_popcount(a) // set bits (add ll)
#define ln '\\n'
#define getcurrtime()                                                          \\
  cerr << "Time = " << ((double)clock() / CLOCKS_PER_SEC) << endl
#define dbgarr(v, s, e)                                                        \\
  cerr << #v << " = ";                                                         \\
  forsn(i, s, e) cerr << v[i] << ", ";                                         \\
  cerr << endl
#define inputfile freopen("input.txt", "r", stdin)
#define outputfile freopen("output.txt", "w", stdout)
#define dbg(args...)                                                           \\
  {                                                                            \\
    string _s = #args;                                                         \\
    replace(_s.begin(), _s.end(), ',', ' ');                                   \\
    stringstream _ss(_s);                                                      \\
    istream_iterator<string> _it(_ss);                                         \\
    err(_it, args);                                                            \\
  }
void err(istream_iterator<string> it) { cerr << endl; }
template <typename T, typename... Args>
void err(istream_iterator<string> it, T a, Args... args) {
  cerr << *it << " = " << a << "\\t";
  err(++it, args...);
}
template <typename T1, typename T2>
ostream &operator<<(ostream &c, pair<T1, T2> &v) {
  c << "(" << v.fi << "," << v.se << ")";
  return c;
}
template <template <class...> class TT, class... T>
ostream &operator<<(ostream &out, TT<T...> &c) {
  out << "{ ";
  forstl(x, c) out << x << " ";
  out << "}";
  return out;
}
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef pair<ll, ll> p64;
typedef pair<int, int> p32;
typedef pair<int, p32> p96;
typedef vector<ll> v64;
typedef vector<int> v32;
typedef vector<v32> vv32;
typedef vector<v64> vv64;
typedef vector<p32> vp32;
typedef vector<p64> vp64;
typedef vector<vp32> vvp32;
typedef map<int, int> m32;
const int LIM = 2e5 + 5, MOD = 998244353;
const ld EPS = 1e-9;

int read() {
  int xx = 0, ff = 1;
  char ch = getchar();
  while (ch > '9' || ch < '0') {
    if (ch == '-')
      ff = -1;
    ch = getchar();
  }
  while (ch >= '0' && ch <= '9') {
    xx = (xx << 3) + (xx << 1) + ch - '0';
    ch = getchar();
  }
  return xx * ff;
}

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

ll add(ll a, ll b) { return (a + b) % MOD; }

ll mul(ll a, ll b) { return (a * b) % MOD; }

ll fastpow(ll a, ll p) {
  if (p == 0)
    return 1;
  ll z = fastpow(a, p / 2);
  z = mul(z, z);
  if (p % 2)
    z = mul(z, a);
  return z;
}

int main() {
  fastio;

  int n, k;
  cin >> n >> k;
  ll ans = 0;
  int a[n];
  forn(i, n) cin >> a[i];

  forn(j, n) {
    forn(i, j) {
      ll rat = mul(k - 1, fastpow(k, MOD - 2));
      rat = fastpow(rat, max(0, j - max(k - 1, i)));
      rat = mul(rat, fastpow(2, MOD - 2));

      if (a[i] < a[j])
        ans = add(ans, rat);
      else {
        ll num = mul(2, fastpow(k, max(0, j - max(k - 1, i)))) -
                 fastpow(k - 1, max(0, j - max(k - 1, i)));
        num %= MOD;
        if (num < 0)
          num += MOD;
        ll denom = mul(2, fastpow(k, max(0, j - max(k - 1, i))));
        num = mul(num, fastpow(denom, MOD - 2));
        num %= MOD;
        ans = add(ans, num);
      }
    }
  }
  cout << ans << ln;

  return 0;
}