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

#define _USE_MATH_DEFINES
#include <algorithm>
#include <assert.h>
#include <boost/multiprecision/cpp_int.hpp>
#include <cctype>
#include <cstdio>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <math.h>
#include <queue>
#include <set>
#include <stack>
#include <string.h>
#include <string>
#include <utility>
#include <vector>
using namespace std;
namespace mp = boost::multiprecision;

#define INF (1 << 20)
#define EPS (1e-10)
#define EQ(a, b) (fabs((a) - (b)) < EPS)

#define ll mp::cpp_int
#define LL mp::cpp_int

#define MAX_N 200000
#define MOD 1000000007ll

int flg[MAX_N + 1];
int p[MAX_N + 1];

// 最大公約数を求める
ll gcd(ll a, ll b) {
  if (a == 0 && b == 0)
    return INF;
  if (a == 0 || b == 0)
    return a + b;
  while (a % b != 0) {
    ll r = a % b;
    a = b;
    b = r;
  }
  return b;
}

// 最小公倍数を求める
ll lcm(ll a, ll b) { return a * b / gcd(a, b); }

int n;
ll period[2 * MAX_N - 1], sum[2 * MAX_N - 1];

// O(n)
// セグメント木の初期化
void init(int n_) {
  // 簡単のため、要素数は２のべき乗にする
  n = 1;
  while (n < n_)
    n *= 2;
  for (int i = 0; i < 2 * MAX_N - 1; ++i)
    period[i] = 1, sum[i] = 0;
}

// O(logn)
// k番目の値(0-origin)を変更
void update(int k, ll per, ll s) {
  // 葉の節点
  k += n - 1;
  period[k] = per;
  sum[k] = s;
  while (k > 0) {
    k = (k - 1) / 2;
    period[k] = lcm(period[k * 2 + 1], period[k * 2 + 2]);
    assert(period[k] >= 0);
    sum[k] = period[k] / period[k * 2 + 1] * sum[k * 2 + 1] +
             period[k] / period[k * 2 + 2] * sum[k * 2 + 2];
    assert(sum[k] >= 0);
  }
}

// O(logn)
//[l, r)のうちの[a, b)の答えを求める。
// kは節点の番号。呼ぶときは根である0を呼ぶこと。
// 例：[0, n)のうち[a, b)の部分の答えを求める：query(a, b, 0, 0, n)
ll query(ll a, ll b, int k, ll l, ll r, ll &per) {
  // cout << "k" << k << endl;
  // cout << "sum[k]" << sum[k] << endl;
  // cout << "period[k]" << period[k] << endl;

  //[l, r) && [a, b) == null    {l r a b} or {a b l r}
  if (r <= a || b <= l)
    return 0;

  //[l, r) && [a, b) == [l, r)    {a l r b}
  if (a <= l && r <= b) {
    per = lcm(per, period[k]);
    return sum[k];
  }

  //{a l b r} or {l a r b} or {l a b r}
  ll vl = query(a, b, k * 2 + 1, l, (l + r) / 2, per);
  ll vr = query(a, b, k * 2 + 2, (l + r) / 2, r, per);
  assert(period[k] % period[k * 2 + 1] == 0 &&
         period[k] % period[k * 2 + 2] == 0);
  return period[k] / period[k * 2 + 1] * vl +
         period[k] / period[k * 2 + 2] * vr;
}

int main() {
  int N, Q;
  cin >> N >> Q;
  init(N);
  for (int i = 0; i < N; ++i) {
    int data;
    cin >> data;
    p[i + 1] = data;
  }
  for (int i = 1; i <= N; ++i) {
    if (flg[i])
      continue;
    flg[i]++;
    int cur = p[i];
    ll s = i;
    ll per = 1;
    while (cur != i) {
      flg[cur]++;
      s += cur;
      ++per;
      cur = p[cur];
    }
    cur = p[i];
    update(i - 1, per, s);
    while (cur != i) {
      update(cur - 1, per, s);
      cur = p[cur];
    }
  }
  for (int i = 0; i < Q; ++i) {
    int l, r;
    cin >> l >> r;
    ll per = 1;
    ll ret = query(l - 1, r, 0, 0, n, per);
    // cout << "ret" << ret << "\\nper" << per << "\\nperiod[0]" << period[0] <<
    // endl;
    assert((ret * per) % period[0] == 0);
    ret = ret * per / period[0];
    cout << ret % MOD << endl;
    l++;
  }
  return 0;
}

Error: TLE

Fixed submission:

#define _USE_MATH_DEFINES
#include <algorithm>
#include <cctype>
#include <cstdio>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <math.h>
#include <queue>
#include <set>
#include <stack>
#include <string.h>
#include <string>
#include <utility>
#include <vector>
using namespace std;

#define INF (1 << 20)
#define EPS (1e-10)
#define EQ(a, b) (fabs((a) - (b)) < EPS)

#define ll long long
#define LL long long

#define MAX_N 100000
#define MOD 1000000007ll

int id[MAX_N + 1];
int p[MAX_N + 1];
int flg[MAX_N + 1];
vector<int> period;
vector<int> sum;

// 最大公約数を求める
ll gcd(ll a, ll b) {
  if (a == 0 && b == 0)
    return INF;
  if (a == 0 || b == 0)
    return a + b;
  while (a % b != 0) {
    ll r = a % b;
    a = b;
    b = r;
  }
  return b;
}

// 最小公倍数を求める
ll lcm(ll a, ll b) { return a * b / gcd(a, b); }

int main() {
  int N, Q;
  cin >> N >> Q;
  for (int i = 0; i < N; ++i) {
    int data;
    cin >> data;
    p[i + 1] = data;
  }
  period.push_back(0); // dummy
  sum.push_back(0);    // dummy
  int cnt = 1;
  for (int i = 1; i <= N; ++i) {
    if (id[i] != 0)
      continue;
    int cur = i;
    int s = 0;
    int per = 0;
    while (id[cur] == 0) {
      id[cur] = cnt;
      s += cur;
      ++per;
      cur = p[cur];
    }
    period.push_back(per);
    sum.push_back(s);
    ++cnt;
  }
  for (int i = 0; i < Q; ++i) {
    int l, r;
    cin >> l >> r;
    vector<int> vec;
    set<int> sett;
    for (; l <= r; ++l) {
      sett.insert(id[l]);
      vec.push_back(id[l]);
    }
    ll per = 1; // 周期
    for_each(sett.begin(), sett.end(),
             [&](int n) { per = lcm(per, period[n]); });
    ll ret = 0;
    for_each(vec.begin(), vec.end(),
             [&](int n) { ret += sum[n] * per / period[n] % MOD; });
    cout << ret << endl;
  }
  return 0;
}