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

#include <bits/stdc++.h>

using namespace std;

#define REP(i, n) for (int i = 0; i < (n); ++i)
#define RREP(i, n) for (int i = (int)(n)-1; i >= 0; --i)
#define FOR(i, a, n) for (int i = (a); i < (n); ++i)
#define RFOR(i, a, n) for (int i = (int)(n)-1; i >= (a); --i)

#define SZ(x) ((int)(x).size())
#define ALL(x) (x).begin(), (x).end()

#define DUMP(x) cerr << #x << " = " << (x) << endl
#define DEBUG(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl;

template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  REP(i, SZ(v)) {
    if (i)
      os << ", ";
    os << v[i];
  }
  return os << "]";
}

template <class T, class U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  return os << "(" << p.first << " " << p.second << ")";
}

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

using ll = long long;
using ull = unsigned long long;
using ld = long double;
using vi = vector<int>;
using vll = vector<ll>;
using vvi = vector<vi>;
using vvll = vector<vll>;

const ll MOD = 1e9 + 7;
const int INF = INT_MAX / 2;
const ll LINF = LLONG_MAX / 2;
const ld eps = 1e-9;

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);

  ll N, K;
  cin >> N >> K;
  vi P(N);
  REP(i, N) cin >> P[i], --P[i];
  vll C(N);
  REP(i, N) cin >> C[i];

  vi used(N);

  ll ans = -LINF;
  REP(i, N) {
    if (used[i])
      continue;
    vi cycle;
    ll sum = 0;
    int j = i;
    do {
      cycle.push_back(j);
      sum += C[j];
      used[j] = 1;
      j = P[j];
    } while (j != i);

    ll size = cycle.size();

    ll rem = K % size;

    ll max0 = -LINF, max1 = -LINF;
    REP(k, size) {
      ll t = 0;
      REP(u, size) {
        t += C[cycle[(k + u) % size]];
        if (u < rem)
          chmax(max0, t);
        else
          chmax(max1, t);
      }
    }
    ll m = K / size;
    if (sum < 0) {
      chmax(ans, max0);
      chmax(ans, max1);
    } else {
      chmax(ans, sum * m + max0);
      if (m > 0)
        chmax(ans, sum * (m - 1) + max1);
    }
  }
  cout << ans << endl;

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

#define REP(i, n) for (int i = 0; i < (n); ++i)
#define RREP(i, n) for (int i = (int)(n)-1; i >= 0; --i)
#define FOR(i, a, n) for (int i = (a); i < (n); ++i)
#define RFOR(i, a, n) for (int i = (int)(n)-1; i >= (a); --i)

#define SZ(x) ((int)(x).size())
#define ALL(x) (x).begin(), (x).end()

#define DUMP(x) cerr << #x << " = " << (x) << endl
#define DEBUG(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl;

template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  REP(i, SZ(v)) {
    if (i)
      os << ", ";
    os << v[i];
  }
  return os << "]";
}

template <class T, class U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  return os << "(" << p.first << " " << p.second << ")";
}

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

using ll = long long;
using ull = unsigned long long;
using ld = long double;
using vi = vector<int>;
using vll = vector<ll>;
using vvi = vector<vi>;
using vvll = vector<vll>;

const ll MOD = 1e9 + 7;
const int INF = INT_MAX / 2;
const ll LINF = LLONG_MAX / 2;
const ld eps = 1e-9;

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);

  int N, K;
  cin >> N >> K;
  vi P(N);
  REP(i, N) cin >> P[i], --P[i];
  vll C(N);
  REP(i, N) cin >> C[i];

  vi used(N);

  ll ans = -INF;
  REP(i, N) {
    if (used[i])
      continue;
    vi cycle;
    ll sum = 0;
    int j = i;
    do {
      cycle.push_back(j);
      sum += C[j];
      used[j] = 1;
      j = P[j];
    } while (j != i);

    ll size = cycle.size();

    {
      ll ma = -INF;
      REP(k, size) {
        ll t = 0;
        REP(u, K % size) {
          t += C[cycle[(k + u) % size]];
          chmax(ma, t);
        }
      }
      ll m = K / size;
      chmax(ans, sum * m + ma);
    }
    {
      ll ma = -INF;
      REP(k, size) {
        ll t = 0;
        REP(u, K) {
          t += C[cycle[(k + u) % size]];
          chmax(ma, t);
        }
      }
      chmax(ans, ma);
    }
  }
  cout << ans << endl;

  return 0;
}