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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <tuple>
#include <vector>

using namespace std;
typedef long long ll;

#define rep(i, n) for (ll i = 0; i < (n); i++)
#define reps(i, f, n) for (ll i = (f); i < (n); i++)
#define repr(i, n) for (ll i = n; i >= 0; i--)
#define repv(v) for (auto it = (v).begin(); it != (v).end(); it++)
#define ALL(x) (x).begin(), (x).end()
#define SZ(x) ((int)(x).size())
#define pb push_back
#define INIT                                                                   \\
  cin.tie(0);                                                                  \\
  ios::sync_with_stdio(false);

template <class T> inline bool chmax(T &a, T b) { return a = (a < b) ? b : a; }
template <class T> inline bool chmin(T &a, T b) { return a = (a > b) ? b : a; }

ll const INF = 1LL << 60;
ll const MOD = 1000000007;
ll const MAX_V = 5005;

struct edge {
  ll to, cost;
};
vector<edge> G[MAX_V];

int main() {
  INIT;

  ll N, K;
  cin >> N >> K;
  vector<ll> P(N), C(N);
  rep(i, N) { cin >> P[i]; }
  rep(i, N) { cin >> C[i]; }

  rep(i, N) {
    P[i]--;
    edge et = {P[i], C[P[i]]};
    G[i].push_back(et);
  }

  ll ans = -INF;
  rep(i, N) {
    ll v = i;
    vector<ll> s;
    ll score = 0;

    rep(n, 5005) {
      edge nv = G[v][0];
      score += nv.cost;
      s.pb(score);

      if (nv.to == i)
        break;
      v = nv.to;
    }

    ll total = s[SZ(s) - 1];
    rep(i, SZ(s)) {
      if (i + 1 > K)
        break;

      if (total > 0) {
        ll e = (K - (i + 1)) / SZ(s);
        s[i] += total * e;
      }

      chmax(ans, s[i]);
    }
  }
  cout << ans << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <tuple>
#include <vector>

using namespace std;
typedef long long ll;

#define rep(i, n) for (ll i = 0; i < (n); i++)
#define reps(i, f, n) for (ll i = (f); i < (n); i++)
#define repr(i, n) for (ll i = n; i >= 0; i--)
#define repv(v) for (auto it = (v).begin(); it != (v).end(); it++)
#define ALL(x) (x).begin(), (x).end()
#define SZ(x) ((int)(x).size())
#define pb push_back
#define INIT                                                                   \\
  cin.tie(0);                                                                  \\
  ios::sync_with_stdio(false);

template <class T> inline bool chmax(T &a, T b) { return a = (a < b) ? b : a; }
template <class T> inline bool chmin(T &a, T b) { return a = (a > b) ? b : a; }

ll const INF = 1LL << 60;
ll const MOD = 1000000007;
ll const MAX_V = 5005;

struct edge {
  ll to, cost;
};
vector<edge> G[MAX_V];
ll N, K;
ll res_dfs;

vector<ll> loop_score;
vector<ll> loop_k;
vector<bool> looped;

ll dfs(ll v, ll p, ll score, ll k) {
  if (k == 0) {
    // cout << p << "->" << v << ": total: " << score << " k: " << k <<
    // endl;
    return score;
  } else {
    for (auto nv : G[v]) {
      // cout << v << "->" << nv.to << ": " << nv.cost << ", score:" <<
      // score
      //      << " k: " << k << endl;

      chmax(res_dfs, score + nv.cost);

      return dfs(nv.to, v, score + nv.cost, k - 1);
    }
  }
}

int main() {
  INIT;

  cin >> N >> K;
  vector<ll> P(N), C(N);
  rep(i, N) { cin >> P[i]; }
  rep(i, N) { cin >> C[i]; }

  rep(i, N) {
    P[i]--;
    edge et = {P[i], C[P[i]]};
    G[i].push_back(et);
  }

  ll ans = -INF;
  rep(i, N) {
    res_dfs = -INF;
    dfs(i, -1, 0, K);
    chmax(ans, res_dfs);
  }
  cout << ans << endl;

  return 0;
}