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

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
const int MOD = 998244353;

int add(int a, int b) {
  a += b;
  if (a >= MOD) {
    return a - MOD;
  }
  if (a < 0) {
    return a + MOD;
  }
  return a;
}

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

int pw(int a, int b) {
  int r = 1;
  while (b) {
    if (b & 1) {
      r = mul(r, a);
    }
    a = mul(a, a);
    b /= 2;
  }
  return r;
}

int dv(int a, int b) { return mul(a, pw(b, MOD - 2)); }

const int N = (int)2e5 + 7;
int n;
int k;
int a[N];
int enter[N];
int kek;
int val[N];
int inv_val[N];
int inv2;
int aib[N];

void clr() {
  for (int i = 1; i <= n; i++) {
    aib[i] = 0;
  }
}

void op(int i, int x) {
  while (i <= n) {
    aib[i] = add(aib[i], x);
    i += i & (-i);
  }
}

int get(int i) {
  int sol = 0;
  while (i) {
    sol = add(sol, aib[i]);
    i -= i & (-i);
  }
  return sol;
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);

  cin >> n >> k;
  inv2 = dv(1, 2);
  kek = dv(k - 1, k);
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
    enter[i] = max(0, i - k);
    val[i] = pw(kek, enter[i]);
    inv_val[i] = dv(1, val[i]);
  }
  int inversion_count = 0;
  for (int i = 1; i <= n; i++) {
    inversion_count = add(inversion_count, i - 1 - get(a[i]));
    op(a[i], 1);
  }
  clr();
  int invs = 0, tot = 0;
  for (int i = n; i >= 1; i--) {
    int x, y;
    y = get(a[i]);
    x = add(tot, -y);
    tot = add(tot, val[i]);
    op(a[i], val[i]);
    x = mul(x, inv2);
    y = mul(y, inv2);
    x = mul(x, inv_val[i]);
    y = mul(y, inv_val[i]);
    invs = add(invs, add(x, -y));
  }
  invs = add(invs, inversion_count);
  cout << invs << "\\n";
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
const int MOD = 998244353;

int add(int a, int b) {
  a += b;
  if (a >= MOD) {
    return a - MOD;
  }
  if (a < 0) {
    return a + MOD;
  }
  return a;
}

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

int pw(int a, int b) {
  int r = 1;
  while (b) {
    if (b & 1) {
      r = mul(r, a);
    }
    a = mul(a, a);
    b /= 2;
  }
  return r;
}

int dv(int a, int b) { return mul(a, pw(b, MOD - 2)); }

const int N = (int)2e5 + 7;
int n;
int k;
int a[N];
int enter[N];
int kek;
int val[N];
int inv_val[N];
int inv2;
int aib[N];

void clr() {
  for (int i = 1; i <= n; i++) {
    aib[i] = 0;
  }
}

void op(int i, int x) {
  while (i <= n) {
    aib[i] += x;
    i += i & (-i);
  }
}

int get(int i) {
  int sol = 0;
  while (i) {
    sol += aib[i];
    i -= i & (-i);
  }
  return sol;
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);

  cin >> n >> k;
  inv2 = dv(1, 2);
  kek = dv(k - 1, k);
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
    enter[i] = max(0, i - k);
    val[i] = pw(kek, enter[i]);
    inv_val[i] = dv(1, val[i]);
  }
  int inversion_count = 0;
  for (int i = 1; i <= n; i++) {
    inversion_count = add(inversion_count, i - 1 - get(a[i]));
    op(a[i], 1);
  }
  clr();
  int invs = 0;
  for (int i = 1; i <= n; i++) {
    int x = 0, y = 0;
    for (int j = i + 1; j <= n; j++) {
      if (a[i] < a[j]) {
        x = add(x, val[j]);
      } else {
        y = add(y, val[j]);
      }
    }
    x = mul(x, inv2);
    y = mul(y, inv2);
    x = mul(x, inv_val[i]);
    y = mul(y, inv_val[i]);
    invs = add(invs, add(x, -y));
  }
  invs = add(invs, inversion_count);
  cout << invs << "\\n";
  return 0;
}