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

#include <bits/stdc++.h>

using namespace std;

namespace {

typedef double real;
typedef long long ll;

template <class T> ostream &operator<<(ostream &os, const vector<T> &vs) {
  if (vs.empty())
    return os << "[]";
  os << "[" << vs[0];
  for (int i = 1; i < vs.size(); i++)
    os << " " << vs[i];
  return os << "]";
}
template <class T> istream &operator>>(istream &is, vector<T> &vs) {
  for (auto it = vs.begin(); it != vs.end(); it++)
    is >> *it;
  return is;
}

const ll MOD = ll(1e9 + 7);

int N, Q;
vector<int> P;
vector<int> L, R;
vector<int> X;
void input() {
  cin >> N >> Q;
  P.clear();
  P.resize(N + 1, -1);
  for (int i = 1; i <= N; i++)
    cin >> P[i];
  L.clear();
  L.resize(Q, -1);
  R.clear();
  R.resize(Q, -1);
  for (int i = 0; i < Q; i++) {
    cin >> L[i] >> R[i];
  }
}

void stupid(int k) {
  X.clear();
  X.resize(N + 1, -1);
  for (int i = 1; i <= N; i++) {
    X[i] = i;
  }
  ll ret = 0;
  while (true) {
    vector<int> Y(N + 1, -1);
    for (int i = 1; i <= N; i++) {
      Y[i] = P[X[i]];
    }
    X = Y;
    for (int i = L[k]; i <= R[k]; i++) {
      ret += X[i];
      ret %= MOD;
    }
    bool flag = true;
    for (int i = L[k]; i <= R[k]; i++) {
      if (X[i] != i) {
        flag = false;
        break;
      }
    }
    if (flag) {
      break;
    }
  }
  cout << ret << endl;
}

ll lcm(ll a, ll b) { return a / __gcd(a, b) * b; }

struct LS {
  ll lcm;
  ll sum;
  LS(ll lcm, ll sum) : lcm(lcm), sum(sum) {}
};
const LS LS_zero = LS(1, 0);
LS operator+(const LS &a, const LS &b) {
  ll x = lcm(a.lcm, b.lcm);
  ll a0 = a.sum * ((x / a.lcm) % MOD) % MOD;
  ll b0 = b.sum * ((x / b.lcm) % MOD) % MOD;
  ll s = (a0 + b0) % MOD;
  return LS(x, s);
}
LS operator+=(LS &a, const LS &b) { return a = a + b; }
ostream &operator<<(ostream &os, const LS &a) {
  return os << "LS(" << a.lcm << "," << a.sum << ")";
}

struct RangeSumQuery {
  int N;
  vector<LS> v;
  static int FindPow2(int x) {
    int x1 = 1;
    while (x1 < x) {
      x1 *= 2;
    }
    return x1;
  }
  RangeSumQuery(int N_) : N(FindPow2(N_)) { v = vector<LS>(N * 2, LS_zero); }
  LS dfs(int s, int t, int a, int b, int k) {
    if (s <= a && b <= t)
      return v[k];
    if (t <= a || b <= s)
      return LS_zero;
    int m = (a + b) / 2;
    return dfs(s, t, a, m, k * 2) + dfs(s, t, m, b, k * 2 + 1);
  }
  // [s, t)
  LS query(int s, int t) { return dfs(s, t, 0, N, 1); }
  void add(int i, LS x) {
    i += N;
    while (i > 0) {
      v[i] += x;
      i /= 2;
    }
  }
  void dump() {
    for (int k = 0;; k++) {
      for (int i = (1 << k); i < (1 << (k + 1)); i++) {
        if (i >= 2 * N) {
          cout << endl;
          return;
        }
        cout << v[i] << '\\t';
      }
      cout << endl;
    }
  }
};

void solve() {
  /* for (int k = 0; k < Q; k++) { stupid(k); } */

  X.clear();
  X.resize(N + 1, -1);
  for (int i = 1; i <= N; i++) {
    X[i] = i;
  }

  vector<ll> C(N + 1, 0);
  C[0] = 1;
  vector<ll> S(N + 1, 0);
  for (int k = 1; k <= N; k++) {
    int o = X[k];
    int x = X[k];
    while (true) {
      x = P[x];
      C[k]++;
      S[k] = (S[k] + x) % MOD;
      if (o == x)
        break;
    }
    assert(C[k] <= 40);
  }

  RangeSumQuery rsq(N + 1);
  for (int i = 1; i <= N; i++) {
    rsq.add(i, LS(C[i], S[i]));
  }

  // rsq.dump();

  for (int q = 0; q < Q; q++) {
    int l = L[q], r = R[q];
    // cerr << l << " " << r << endl;
    LS x = rsq.query(l, r + 1);
    cout << x.sum << endl;
    // cout << x << endl;
  }
}
} // namespace

int main() {
  input();
  solve();
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

namespace {

typedef double real;
typedef long long ll;

template <class T> ostream &operator<<(ostream &os, const vector<T> &vs) {
  if (vs.empty())
    return os << "[]";
  os << "[" << vs[0];
  for (int i = 1; i < vs.size(); i++)
    os << " " << vs[i];
  return os << "]";
}
template <class T> istream &operator>>(istream &is, vector<T> &vs) {
  for (auto it = vs.begin(); it != vs.end(); it++)
    is >> *it;
  return is;
}

const ll MOD = ll(1e9 + 7);

int N, Q;
vector<int> P;
vector<int> L, R;
vector<int> X;
void input() {
  cin >> N >> Q;
  P.clear();
  P.resize(N + 1, -1);
  for (int i = 1; i <= N; i++)
    cin >> P[i];
  L.clear();
  L.resize(Q, -1);
  R.clear();
  R.resize(Q, -1);
  for (int i = 0; i < Q; i++) {
    cin >> L[i] >> R[i];
  }
}

void stupid(int k) {
  X.clear();
  X.resize(N + 1, -1);
  for (int i = 1; i <= N; i++) {
    X[i] = i;
  }
  ll ret = 0;
  while (true) {
    vector<int> Y(N + 1, -1);
    for (int i = 1; i <= N; i++) {
      Y[i] = P[X[i]];
    }
    X = Y;
    for (int i = L[k]; i <= R[k]; i++) {
      ret += X[i];
      ret %= MOD;
    }
    bool flag = true;
    for (int i = L[k]; i <= R[k]; i++) {
      if (X[i] != i) {
        flag = false;
        break;
      }
    }
    if (flag) {
      break;
    }
  }
  cout << ret << endl;
}

ll lcm(ll a, ll b) { return a / __gcd(a, b) * b; }

void solve() {
  /* for (int k = 0; k < Q; k++) { stupid(k); } */

  X.clear();
  X.resize(N + 1, -1);
  for (int i = 1; i <= N; i++) {
    X[i] = i;
  }

  vector<int> C(N + 1, 0);
  vector<int> S(N + 1, 0);
  for (int k = 1; k <= N; k++) {
    int o = X[k];
    int x = X[k];
    while (true) {
      x = P[x];
      C[k]++;
      S[k] = (S[k] + x) % MOD;
      if (o == x)
        break;
    }
    assert(C[k] <= 40);
  }

  for (int q = 0; q < Q; q++) {
    int l = L[q], r = R[q];
    ll x = 1;
    for (int i = l; i <= r; i++) {
      x = lcm(x, C[i]);
    }
    ll s = 0;
    for (int i = l; i <= r; i++) {
      s = (s + S[i] * ((x / C[i]) % MOD)) % MOD;
    }
    cout << s << endl;
  }
}
} // namespace

int main() {
  input();
  solve();
  return 0;
}