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

#include <algorithm>
#include <array>
#include <cassert>
#include <iostream>
#include <vector>
using namespace std;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << '{';
  string sep;
  for (const auto &x : v)
    os << sep << x, sep = ", ";
  return os << '}';
}
template <typename T, size_t size>
ostream &operator<<(ostream &os, const array<T, size> &arr) {
  os << '{';
  string sep;
  for (const auto &x : arr)
    os << sep << x, sep = ", ";
  return os << '}';
}
template <typename A, typename B>
ostream &operator<<(ostream &os, const pair<A, B> &p) {
  return os << '(' << p.first << ", " << p.second << ')';
}

void dbg_out() { cerr << endl; }
template <typename Head, typename... Tail> void dbg_out(Head H, Tail... T) {
  cerr << ' ' << H;
  dbg_out(T...);
}

#ifdef NEAL_DEBUG
#define dbg(...) cerr << "(" << #__VA_ARGS__ << "):", dbg_out(__VA_ARGS__)
#else
#define dbg(...)
#endif

template <const int &MOD> struct _m_int {
  int val;

  _m_int(int64_t v = 0) {
    if (v < 0)
      v = v % MOD + MOD;
    if (v >= MOD)
      v %= MOD;
    val = int(v);
  }

  _m_int(uint64_t v) {
    if (v >= MOD)
      v %= MOD;
    val = int(v);
  }

  _m_int(int v) : _m_int(int64_t(v)) {}
  _m_int(unsigned v) : _m_int(uint64_t(v)) {}

  static int inv_mod(int a, int m = MOD) {
    // https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm#Example
    int g = m, r = a, x = 0, y = 1;

    while (r != 0) {
      int q = g / r;
      g %= r;
      swap(g, r);
      x -= q * y;
      swap(x, y);
    }

    return x < 0 ? x + m : x;
  }

  explicit operator int() const { return val; }
  explicit operator unsigned() const { return val; }
  explicit operator int64_t() const { return val; }
  explicit operator uint64_t() const { return val; }
  explicit operator double() const { return val; }
  explicit operator long double() const { return val; }

  _m_int &operator+=(const _m_int &other) {
    val -= MOD - other.val;
    if (val < 0)
      val += MOD;
    return *this;
  }

  _m_int &operator-=(const _m_int &other) {
    val -= other.val;
    if (val < 0)
      val += MOD;
    return *this;
  }

  static unsigned fast_mod(uint64_t x, unsigned m = MOD) {
#if !defined(_WIN32) || defined(_WIN64)
    return unsigned(x % m);
#endif
    // Optimized mod for Codeforces 32-bit machines.
    // x must be less than 2^32 * m for this to work, so that x / m fits in an
    // unsigned 32-bit int.
    unsigned x_high = unsigned(x >> 32), x_low = unsigned(x);
    unsigned quot, rem;
    asm("divl %4\\n" : "=a"(quot), "=d"(rem) : "d"(x_high), "a"(x_low), "r"(m));
    return rem;
  }

  _m_int &operator*=(const _m_int &other) {
    val = fast_mod(uint64_t(val) * other.val);
    return *this;
  }

  _m_int &operator/=(const _m_int &other) { return *this *= other.inv(); }

  friend _m_int operator+(const _m_int &a, const _m_int &b) {
    return _m_int(a) += b;
  }
  friend _m_int operator-(const _m_int &a, const _m_int &b) {
    return _m_int(a) -= b;
  }
  friend _m_int operator*(const _m_int &a, const _m_int &b) {
    return _m_int(a) *= b;
  }
  friend _m_int operator/(const _m_int &a, const _m_int &b) {
    return _m_int(a) /= b;
  }

  _m_int &operator++() {
    val = val == MOD - 1 ? 0 : val + 1;
    return *this;
  }

  _m_int &operator--() {
    val = val == 0 ? MOD - 1 : val - 1;
    return *this;
  }

  _m_int operator++(int) {
    _m_int before = *this;
    ++*this;
    return before;
  }
  _m_int operator--(int) {
    _m_int before = *this;
    --*this;
    return before;
  }

  _m_int operator-() const { return val == 0 ? 0 : MOD - val; }

  friend bool operator==(const _m_int &a, const _m_int &b) {
    return a.val == b.val;
  }
  friend bool operator!=(const _m_int &a, const _m_int &b) {
    return a.val != b.val;
  }
  friend bool operator<(const _m_int &a, const _m_int &b) {
    return a.val < b.val;
  }
  friend bool operator>(const _m_int &a, const _m_int &b) {
    return a.val > b.val;
  }
  friend bool operator<=(const _m_int &a, const _m_int &b) {
    return a.val <= b.val;
  }
  friend bool operator>=(const _m_int &a, const _m_int &b) {
    return a.val >= b.val;
  }

  _m_int inv() const { return inv_mod(val); }

  _m_int pow(int64_t p) const {
    if (p < 0)
      return inv().pow(-p);

    _m_int a = *this, result = 1;

    while (p > 0) {
      if (p & 1)
        result *= a;

      p >>= 1;

      if (p > 0)
        a *= a;
    }

    return result;
  }

  friend ostream &operator<<(ostream &os, const _m_int &m) {
    return os << m.val;
  }
};

extern const int MOD = 998244353;
using mod_int = _m_int<MOD>;

const mod_int INV2 = mod_int(2).inv();

mod_int INV1H, INV1H2;

int N, K;
mod_int inv_K, H;
vector<mod_int> H_pow;

struct segment_change {
  mod_int h_add;

  segment_change(mod_int _h_add = 0) : h_add(_h_add) {}

  void reset() { h_add = 0; }

  bool has_change() const { return h_add != 0; }

  // Return the combined result of applying this segment change followed by
  // `other`.
  segment_change combine(const segment_change &other) const {
    return segment_change(h_add + other.h_add);
  }
};

struct segment {
  mod_int sum, h_sum;
  int length = 0;

  segment(mod_int _sum = 0, mod_int _h_sum = 0, int _length = 0)
      : sum(_sum), h_sum(_h_sum), length(_length) {}

  void apply(int _length, const segment_change &change) {
    // dbg(length, _length);
    assert(length <= _length);

    if (change.h_add != 0) {
      sum += change.h_add * (1 - H_pow[length]) * INV1H;
      h_sum += change.h_add * (1 - H_pow[2 * length]) * INV1H2;
    }
  }

  void join(const segment &other) {
    sum += other.sum;
    h_sum += H_pow[length] * other.h_sum;
    length += other.length;
  }

  void join(const segment &a, const segment &b) {
    *this = a;
    join(b);
  }
};

pair<int, int> right_half[32];

struct seg_tree {
  int tree_n = 0;
  vector<segment> tree;
  vector<segment_change> changes;

  seg_tree(int n = -1) {
    if (n >= 0)
      init(n);
  }

  void init(int n) {
    tree_n = 1;

    while (tree_n < n)
      tree_n *= 2;

    tree.assign(2 * tree_n, segment());
    changes.assign(tree_n, segment_change());
  }

  // Builds our tree from an array in O(n).
  void build(const vector<segment> &initial) {
    int n = int(initial.size());
    init(n);
    assert(n <= tree_n);

    for (int i = 0; i < n; i++)
      tree[tree_n + i] = initial[i];

    for (int position = tree_n - 1; position > 0; position--)
      tree[position].join(tree[2 * position], tree[2 * position + 1]);
  }

  void apply_and_combine(int position, int length,
                         const segment_change &change) {
    tree[position].apply(length, change);
    if (position < tree_n)
      changes[position] = changes[position].combine(change);
  }

  void push_down(int position, int length) {
    if (changes[position].has_change()) {
      apply_and_combine(2 * position, length / 2, changes[position]);
      changes[position].h_add *= H_pow[length / 2];
      apply_and_combine(2 * position + 1, length / 2, changes[position]);
      changes[position].reset();
    }
  }

  // Calls push_down for all necessary nodes in order to query the range [a, b).
  void push_all(int a, int b) {
    assert(0 <= a && a < b && b <= tree_n);
    a += tree_n;
    b += tree_n - 1;

    for (int up = 31 - __builtin_clz(tree_n); up > 0; up--) {
      int x = a >> up, y = b >> up;
      push_down(x, 1 << up);
      if (x != y)
        push_down(y, 1 << up);
    }
  }

  void join_and_apply(int position, int length) {
    tree[position].join(tree[2 * position], tree[2 * position + 1]);
    tree[position].apply(length, changes[position]);
  }

  // Calls join for all necessary nodes after updating the range [a, b).
  void join_all(int a, int b) {
    assert(0 <= a && a < b && b <= tree_n);
    a += tree_n;
    b += tree_n - 1;
    int length = 1;

    while (a > 1) {
      a /= 2;
      b /= 2;
      length *= 2;
      join_and_apply(a, length);
      if (a != b)
        join_and_apply(b, length);
    }
  }

  template <typename T_range_op>
  void process_range(int a, int b, bool needs_join, T_range_op &&range_op) {
    if (a == b)
      return;
    push_all(a, b);
    int original_a = a, original_b = b;
    int length = 1, r_size = 0;

    for (a += tree_n, b += tree_n; a < b; a /= 2, b /= 2, length *= 2) {
      if (a & 1)
        range_op(a++, length);

      if (b & 1)
        right_half[r_size++] = {--b, length};
    }

    for (int i = r_size - 1; i >= 0; i--)
      range_op(right_half[i].first, right_half[i].second);

    if (needs_join)
      join_all(original_a, original_b);
  }

  segment query(int a, int b) {
    assert(0 <= a && a <= b && b <= tree_n);
    segment answer;

    process_range(a, b, false,
                  [&](int position, int) { answer.join(tree[position]); });

    return answer;
  }

  void update(int a, int b, const segment_change &change) {
    assert(0 <= a && a <= b && b <= tree_n);

    process_range(a, b, true, [&](int position, int length) {
      int start = length * (position - (1 << (31 - __builtin_clz(position))));
      apply_and_combine(position, length,
                        segment_change(change.h_add * H_pow[start - a]));
    });
  }

  vector<segment> to_array() {
    for (int i = 1; i < tree_n; i++)
      push_down(i, tree_n >> (31 - __builtin_clz(i)));

    vector<segment> segs(tree_n);

    for (int i = 0; i < tree_n; i++)
      segs[i] = tree[tree_n + i];

    return segs;
  }
};

int main() {
  ios::sync_with_stdio(false);
#ifndef NEAL_DEBUG
  cin.tie(nullptr);
#endif

  cin >> N >> K;
  inv_K = mod_int(K).inv();
  H = 1 - inv_K;
  H_pow.assign(10 * N + 1, 1);

  for (int i = 1; i <= 10 * N; i++)
    H_pow[i] = H * H_pow[i - 1];

  INV1H = 1 / (1 - H);
  INV1H2 = 1 / (1 - H * H);

  vector<int> P(N);
  vector<int> location(N + 1, -1);

  for (int i = 0; i < N; i++) {
    cin >> P[i];
    location[P[i]] = i;
  }

  seg_tree tree(N - K);
  tree.build(vector<segment>(N - K, segment(0, 0, 1)));
  mod_int last_K = 0;

  auto add = [&](int start) {
    start -= K - 1;
    start = max(start, 0);
    assert(0 <= start && start <= N - K);
    tree.update(start, N - K, segment_change(inv_K));
    last_K += inv_K * H_pow[N - K - start];
  };

  auto query_sum = [&](int start) {
    start -= K - 1;
    start = max(start, 0);
    assert(0 <= start && start <= N - K);
    mod_int sum = tree.query(0, start).sum;
    sum += H * tree.query(start, N - K).h_sum;
    sum += last_K * H_pow[N - K - start + 1] * INV2 * K;
    return sum;
  };

  auto query_here = [&](int start) {
    start -= K - 1;
    start = max(start, 0);
    assert(0 <= start && start <= N - K);
    mod_int sum = inv_K * tree.query(start, N - K).h_sum;
    sum += last_K * K * H_pow[N - K - start] * inv_K;
    return sum;
  };

  mod_int answer = 0;

  for (int x = N; x > 0; x--) {
    int start = location[x];
    // dbg(start, query_sum(start), query_here(start) * INV2);
    answer += query_sum(start);
    answer += query_here(start) * INV2;
    add(start);

    // auto arr = tree.to_array();
    // for (int i = 0; i < N - K; i++)
    //     cerr << arr[i].sum << ' ';
    // cerr << last_K << endl;
  }

  cout << answer << '\\n';
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <array>
#include <cassert>
#include <iostream>
#include <vector>
using namespace std;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << '{';
  string sep;
  for (const auto &x : v)
    os << sep << x, sep = ", ";
  return os << '}';
}
template <typename T, size_t size>
ostream &operator<<(ostream &os, const array<T, size> &arr) {
  os << '{';
  string sep;
  for (const auto &x : arr)
    os << sep << x, sep = ", ";
  return os << '}';
}
template <typename A, typename B>
ostream &operator<<(ostream &os, const pair<A, B> &p) {
  return os << '(' << p.first << ", " << p.second << ')';
}

void dbg_out() { cerr << endl; }
template <typename Head, typename... Tail> void dbg_out(Head H, Tail... T) {
  cerr << ' ' << H;
  dbg_out(T...);
}

#ifdef NEAL_DEBUG
#define dbg(...) cerr << "(" << #__VA_ARGS__ << "):", dbg_out(__VA_ARGS__)
#else
#define dbg(...)
#endif

template <const int &MOD> struct _m_int {
  int val;

  _m_int(int64_t v = 0) {
    if (v < 0)
      v = v % MOD + MOD;
    if (v >= MOD)
      v %= MOD;
    val = int(v);
  }

  _m_int(uint64_t v) {
    if (v >= MOD)
      v %= MOD;
    val = int(v);
  }

  _m_int(int v) : _m_int(int64_t(v)) {}
  _m_int(unsigned v) : _m_int(uint64_t(v)) {}

  static int inv_mod(int a, int m = MOD) {
    // https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm#Example
    int g = m, r = a, x = 0, y = 1;

    while (r != 0) {
      int q = g / r;
      g %= r;
      swap(g, r);
      x -= q * y;
      swap(x, y);
    }

    return x < 0 ? x + m : x;
  }

  explicit operator int() const { return val; }
  explicit operator unsigned() const { return val; }
  explicit operator int64_t() const { return val; }
  explicit operator uint64_t() const { return val; }
  explicit operator double() const { return val; }
  explicit operator long double() const { return val; }

  _m_int &operator+=(const _m_int &other) {
    val -= MOD - other.val;
    if (val < 0)
      val += MOD;
    return *this;
  }

  _m_int &operator-=(const _m_int &other) {
    val -= other.val;
    if (val < 0)
      val += MOD;
    return *this;
  }

  static unsigned fast_mod(uint64_t x, unsigned m = MOD) {
#if !defined(_WIN32) || defined(_WIN64)
    return unsigned(x % m);
#endif
    // Optimized mod for Codeforces 32-bit machines.
    // x must be less than 2^32 * m for this to work, so that x / m fits in an
    // unsigned 32-bit int.
    unsigned x_high = unsigned(x >> 32), x_low = unsigned(x);
    unsigned quot, rem;
    asm("divl %4\\n" : "=a"(quot), "=d"(rem) : "d"(x_high), "a"(x_low), "r"(m));
    return rem;
  }

  _m_int &operator*=(const _m_int &other) {
    val = fast_mod(uint64_t(val) * other.val);
    return *this;
  }

  _m_int &operator/=(const _m_int &other) { return *this *= other.inv(); }

  friend _m_int operator+(const _m_int &a, const _m_int &b) {
    return _m_int(a) += b;
  }
  friend _m_int operator-(const _m_int &a, const _m_int &b) {
    return _m_int(a) -= b;
  }
  friend _m_int operator*(const _m_int &a, const _m_int &b) {
    return _m_int(a) *= b;
  }
  friend _m_int operator/(const _m_int &a, const _m_int &b) {
    return _m_int(a) /= b;
  }

  _m_int &operator++() {
    val = val == MOD - 1 ? 0 : val + 1;
    return *this;
  }

  _m_int &operator--() {
    val = val == 0 ? MOD - 1 : val - 1;
    return *this;
  }

  _m_int operator++(int) {
    _m_int before = *this;
    ++*this;
    return before;
  }
  _m_int operator--(int) {
    _m_int before = *this;
    --*this;
    return before;
  }

  _m_int operator-() const { return val == 0 ? 0 : MOD - val; }

  friend bool operator==(const _m_int &a, const _m_int &b) {
    return a.val == b.val;
  }
  friend bool operator!=(const _m_int &a, const _m_int &b) {
    return a.val != b.val;
  }
  friend bool operator<(const _m_int &a, const _m_int &b) {
    return a.val < b.val;
  }
  friend bool operator>(const _m_int &a, const _m_int &b) {
    return a.val > b.val;
  }
  friend bool operator<=(const _m_int &a, const _m_int &b) {
    return a.val <= b.val;
  }
  friend bool operator>=(const _m_int &a, const _m_int &b) {
    return a.val >= b.val;
  }

  _m_int inv() const { return inv_mod(val); }

  _m_int pow(int64_t p) const {
    if (p < 0)
      return inv().pow(-p);

    _m_int a = *this, result = 1;

    while (p > 0) {
      if (p & 1)
        result *= a;

      p >>= 1;

      if (p > 0)
        a *= a;
    }

    return result;
  }

  friend ostream &operator<<(ostream &os, const _m_int &m) {
    return os << m.val;
  }
};

extern const int MOD = 998244353;
using mod_int = _m_int<MOD>;

int N, K;
mod_int inv_K, H;
vector<mod_int> H_pow;

struct segment_change {
  mod_int h_add;

  segment_change(mod_int _h_add = 0) : h_add(_h_add) {}

  void reset() { h_add = 0; }

  bool has_change() const { return h_add != 0; }

  // Return the combined result of applying this segment change followed by
  // `other`.
  segment_change combine(const segment_change &other) const {
    return segment_change(h_add + other.h_add);
  }
};

struct segment {
  mod_int sum, h_sum;
  int length = 0;

  segment(mod_int _sum = 0, mod_int _h_sum = 0, int _length = 0)
      : sum(_sum), h_sum(_h_sum), length(_length) {}

  void apply(int _length, const segment_change &change) {
    // dbg(length, _length);
    assert(length <= _length);

    if (change.h_add != 0) {
      sum += change.h_add * (1 - H_pow[length]) / (1 - H);
      h_sum += change.h_add * (1 - H_pow[2 * length]) / (1 - H * H);
    }
  }

  void join(const segment &other) {
    sum += other.sum;
    h_sum += H_pow[length] * other.h_sum;
    length += other.length;
  }

  void join(const segment &a, const segment &b) {
    *this = a;
    join(b);
  }
};

pair<int, int> right_half[32];

struct seg_tree {
  int tree_n = 0;
  vector<segment> tree;
  vector<segment_change> changes;

  seg_tree(int n = -1) {
    if (n >= 0)
      init(n);
  }

  void init(int n) {
    tree_n = 1;

    while (tree_n < n)
      tree_n *= 2;

    tree.assign(2 * tree_n, segment());
    changes.assign(tree_n, segment_change());
  }

  // Builds our tree from an array in O(n).
  void build(const vector<segment> &initial) {
    int n = int(initial.size());
    init(n);
    assert(n <= tree_n);

    for (int i = 0; i < n; i++)
      tree[tree_n + i] = initial[i];

    for (int position = tree_n - 1; position > 0; position--)
      tree[position].join(tree[2 * position], tree[2 * position + 1]);
  }

  void apply_and_combine(int position, int length,
                         const segment_change &change) {
    tree[position].apply(length, change);
    if (position < tree_n)
      changes[position] = changes[position].combine(change);
  }

  void push_down(int position, int length) {
    if (changes[position].has_change()) {
      apply_and_combine(2 * position, length / 2, changes[position]);
      changes[position].h_add *= H_pow[length / 2];
      apply_and_combine(2 * position + 1, length / 2, changes[position]);
      changes[position].reset();
    }
  }

  // Calls push_down for all necessary nodes in order to query the range [a, b).
  void push_all(int a, int b) {
    assert(0 <= a && a < b && b <= tree_n);
    a += tree_n;
    b += tree_n - 1;

    for (int up = 31 - __builtin_clz(tree_n); up > 0; up--) {
      int x = a >> up, y = b >> up;
      push_down(x, 1 << up);
      if (x != y)
        push_down(y, 1 << up);
    }
  }

  void join_and_apply(int position, int length) {
    tree[position].join(tree[2 * position], tree[2 * position + 1]);
    tree[position].apply(length, changes[position]);
  }

  // Calls join for all necessary nodes after updating the range [a, b).
  void join_all(int a, int b) {
    assert(0 <= a && a < b && b <= tree_n);
    a += tree_n;
    b += tree_n - 1;
    int length = 1;

    while (a > 1) {
      a /= 2;
      b /= 2;
      length *= 2;
      join_and_apply(a, length);
      if (a != b)
        join_and_apply(b, length);
    }
  }

  template <typename T_range_op>
  void process_range(int a, int b, bool needs_join, T_range_op &&range_op) {
    if (a == b)
      return;
    push_all(a, b);
    int original_a = a, original_b = b;
    int length = 1, r_size = 0;

    for (a += tree_n, b += tree_n; a < b; a /= 2, b /= 2, length *= 2) {
      if (a & 1)
        range_op(a++, length);

      if (b & 1)
        right_half[r_size++] = {--b, length};
    }

    for (int i = r_size - 1; i >= 0; i--)
      range_op(right_half[i].first, right_half[i].second);

    if (needs_join)
      join_all(original_a, original_b);
  }

  segment query(int a, int b) {
    assert(0 <= a && a <= b && b <= tree_n);
    segment answer;

    process_range(a, b, false,
                  [&](int position, int) { answer.join(tree[position]); });

    return answer;
  }

  void update(int a, int b, const segment_change &change) {
    assert(0 <= a && a <= b && b <= tree_n);

    process_range(a, b, true, [&](int position, int length) {
      int start = length * (position - (1 << (31 - __builtin_clz(position))));
      apply_and_combine(position, length,
                        segment_change(change.h_add * H_pow[start - a]));
    });
  }

  vector<segment> to_array() {
    for (int i = 1; i < tree_n; i++)
      push_down(i, tree_n >> (31 - __builtin_clz(i)));

    vector<segment> segs(tree_n);

    for (int i = 0; i < tree_n; i++)
      segs[i] = tree[tree_n + i];

    return segs;
  }
};

int main() {
  ios::sync_with_stdio(false);
#ifndef NEAL_DEBUG
  cin.tie(nullptr);
#endif

  cin >> N >> K;
  inv_K = mod_int(K).inv();
  H = 1 - inv_K;
  H_pow.assign(10 * N + 1, 1);

  for (int i = 1; i <= 10 * N; i++)
    H_pow[i] = H * H_pow[i - 1];

  vector<int> P(N);
  vector<int> location(N + 1, -1);

  for (int i = 0; i < N; i++) {
    cin >> P[i];
    location[P[i]] = i;
  }

  seg_tree tree(N - K);
  tree.build(vector<segment>(N - K, segment(0, 0, 1)));
  mod_int last_K = 0;

  auto add = [&](int start) {
    start -= K - 1;
    start = max(start, 0);
    assert(0 <= start && start <= N - K);
    tree.update(start, N - K, segment_change(inv_K));
    last_K += inv_K * H_pow[N - K - start];
  };

  auto query_sum = [&](int start) {
    start -= K - 1;
    start = max(start, 0);
    assert(0 <= start && start <= N - K);
    mod_int sum = tree.query(0, start).sum;
    sum += H * tree.query(start, N - K).h_sum;
    sum += last_K * H_pow[N - K - start + 1] / 2 * K;
    return sum;
  };

  auto query_here = [&](int start) {
    start -= K - 1;
    start = max(start, 0);
    assert(0 <= start && start <= N - K);
    mod_int sum = inv_K * tree.query(start, N - K).h_sum;
    sum += last_K * K * H_pow[N - K - start] * inv_K;
    return sum;
  };

  mod_int answer = 0;

  for (int x = N; x > 0; x--) {
    int start = location[x];
    // dbg(start, query_sum(start), query_here(start) / 2);
    answer += query_sum(start);
    answer += query_here(start) / 2;
    add(start);

    // auto arr = tree.to_array();
    // for (int i = 0; i < N - K; i++)
    //     cerr << arr[i].sum << ' ';
    // cerr << last_K << endl;
  }

  cout << answer << '\\n';
}