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

#include <bits/stdc++.h>
using namespace std;

using ll = long long;
using pii = pair<int, int>;
template <class T> using V = vector<T>;
template <class T> using VV = V<V<T>>;

#define pb push_back
#define eb emplace_back
#define mp make_pair
#define fi first
#define se second
#define rep(i, n) rep2(i, 0, n)
#define rep2(i, m, n) for (int i = m; i < (n); i++)
#define per(i, b) per2(i, 0, b)
#define per2(i, a, b) for (int i = int(b) - 1; i >= int(a); i--)
#define ALL(c) (c).begin(), (c).end()
#define SZ(x) ((int)(x).size())

constexpr ll TEN(int n) { return (n == 0) ? 1 : 10 * TEN(n - 1); }

template <class T, class U> void chmin(T &t, const U &u) {
  if (t > u)
    t = u;
}
template <class T, class U> void chmax(T &t, const U &u) {
  if (t < u)
    t = u;
}

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

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

#ifdef LOCAL
void debug_out() { cerr << endl; }
template <typename Head, typename... Tail> void debug_out(Head H, Tail... T) {
  cerr << " " << H;
  debug_out(T...);
}
#define debug(...)                                                             \\
  cerr << __LINE__ << " [" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#define dump(x) cerr << __LINE__ << " " << #x << " = " << (x) << endl
#else
#define debug(...) (void(0))
#define dump(x) (void(0))
#endif

constexpr ll INF = TEN(18);

int main() {
  int N, K;
  cin >> N >> K;
  V<int> P(N), C(N);
  rep(i, N) cin >> P[i], --P[i];
  rep(i, N) cin >> C[i];

  ll ans = -INF;

  rep(i, N) {
    int v = i;
    ll cycle_sum = 0;
    int cycle_cnt = 0;

    while (true) {
      cycle_cnt++;
      cycle_sum += C[v];
      v = P[v];
      if (v == i)
        break;
    }

    ll path = 0;
    int cnt = 0;

    while (true) {
      cnt++;
      path += C[v];

      if (cnt > K)
        break;

      int num = (K - cnt) / cycle_cnt;
      ll score = path + max(0ll, cycle_sum) * num;
      chmax(ans, score);

      v = P[v];
      if (v == i)
        break;
    }
  }

  cout << ans << endl;

  return 0;
}


Error: TLE

Fixed submission:


// Problem : D - Moving Piece
// Contest : AtCoder - AtCoder Beginner Contest 175
// URL : https://atcoder.jp/contests/abc175/tasks/abc175_d
// Memory Limit : 1024 MB
// Time Limit : 3000 ms
// Powered by CP Editor (https://github.com/cpeditor/cpeditor)

/*

WINNERS NEVER QUIT AND QUITTERS NEVER WIN!!

*/

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef pair<long, long> pl;
typedef pair<ll, ll> pll;
typedef vector<long> vl;
typedef vector<bool> vb;
typedef vector<ll> vll;
typedef vector<vl> vvl;
typedef vector<vb> vvb;
typedef vector<vll> vvll;
typedef vector<pll> vpll;
typedef vector<string> vs;
typedef unordered_map<ll, ll> umll;

#define FOR(i, a, b) for (long long i = a; i < b; ++i)
#define REV(i, a, b) for (long long i = a; i >= b; i--)
#define F first
#define S second
#define pb push_back
#define mp make_pair
#define ub upper_bound
#define lb lower_bound
#define all(v) v.begin(), v.end()
#define rall(v) v.rbegin(), v.rend()
#define tc                                                                     \\
  ll tests;                                                                    \\
  cin >> tests;                                                                \\
  while (tests--)
#define io                                                                     \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL);                                                               \\
  cout.tie(NULL)
#define coutv(v)                                                               \\
  for (auto it : (v))                                                          \\
    cout << it << " ";                                                         \\
  newl;
#define cout2d(v)                                                              \\
  for (auto it : (v)) {                                                        \\
    for (auto j : it)                                                          \\
      cout << j << " ";                                                        \\
    newl;                                                                      \\
  }
#define cinv(v, n)                                                             \\
  vll(v)(n);                                                                   \\
  FOR(i, 0, (n)) { cin >> v[i]; }
#define cinvg(v, n)                                                            \\
  (v).resize(n);                                                               \\
  FOR(i, 0, (n)) { cin >> v[i]; }
#define cin2d(v, n, m)                                                         \\
  vvll(v)(n, vll(m, 0));                                                       \\
  FOR(i, 0, n) {                                                               \\
    FOR(j, 0, m) { cin >> v[i][j]; }                                           \\
  }
#define cin2dg(v, n, m)                                                        \\
  (v).resize(n, vll(m));                                                       \\
  FOR(i, 0, n) {                                                               \\
    FOR(j, 0, m) { cin >> v[i][j]; }                                           \\
  }
#define newl cout << "\\n"
#define mod 1000000007
#define INF LLONG_MAX / 2

ll n, k, ans = -INF;
vll p, c;

int main() {
  io;
  cin >> n >> k;
  cinvg(p, n) cinvg(c, n) FOR(i, 0, n) { p[i]--; }

  FOR(i, 0, n) {
    ll temp = 0, x = i;
    FOR(j, 0, k) {
      temp += c[p[x]];
      x = p[x];
      ans = max(ans, temp);
    }
  }
  cout << ans;

  return 0;
}