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

#include <bits/stdc++.h>

#define rep2(x, fr, to) for (int x = (fr); x < (to); x++)
#define rep(x, to) for (int x = 0; x < (to); x++)
#define repr(x, fr, to) for (int x = (fr); x >= (to); x--)
#define all(c) c.begin(), c.end()
#define sz(v) (int)v.size()

using namespace std;
typedef long long ll;
typedef vector<int> VI;
typedef pair<int, int> pii;
typedef vector<ll> VL;
const int MD = 1e9 + 7;
void dbg() { cerr << "\\n"; }
template <class F, class... S> void dbg(const F &f, const S &...s) {
  cerr << f << ": ";
  dbg(s...);
}

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  VI p(n), c(n);
  rep(i, n) {
    cin >> p[i];
    p[i]--;
  }
  rep(i, n) cin >> c[i];

  vector<VI> hsu(n, VI(n));
  vector<VL> pt(n, VL(n));
  VI vc(n, -1);
  rep(i, n) {
    ll cr = 1, cr2 = 0;
    int z = p[i];
    cr2 = c[z];
    if (vc[z] == -1)
      vc[z] = i;
    hsu[i][z] = cr;
    pt[i][z] = cr2;
    rep(j, n) {
      z = p[z];
      cr++;
      cr2 += c[z];
      if (hsu[i][z] != 0)
        break;
      hsu[i][z] = cr;
      pt[i][z] = cr2;
      if (vc[z] == -1)
        vc[z] = i;
    }
  }

  VI lp(n);
  VL sk(n + 3);
  rep(i, n) {
    lp[vc[i]]++;
    sk[vc[i]] += c[i];
  }
  // for(auto xl: hsu){for(auto x :xl) cerr<<x<<" ";dbg("hsu");}
  // for(auto xl: pt){for(auto x :xl) cerr<<x<<" ";dbg("pt");}
  // for(auto x :sk) cerr<<x<<" "; dbg("sk");
  // for(auto x :vc) cerr<<x<<" "; dbg("lp");
  ll ans = -1LL << 55;
  rep(s, n) rep(e, n) {
    if (hsu[s][e] > k || hsu[s][e] == 0)
      continue;
    ll cnd = pt[s][e];
    ll z = (k - hsu[s][e]) / lp[vc[e]];
    cnd += max(0LL, sk[vc[e]]) * z;
    ans = max(cnd, ans);
  }
  cout << ans << "\\n";

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

#define rep2(x, fr, to) for (int x = (fr); x < (to); x++)
#define rep(x, to) for (int x = 0; x < (to); x++)
#define repr(x, fr, to) for (int x = (fr); x >= (to); x--)
#define all(c) c.begin(), c.end()
#define sz(v) (int)v.size()

using namespace std;
typedef long long ll;
typedef vector<int> VI;
typedef pair<int, int> pii;
typedef vector<ll> VL;
const int MD = 1e9 + 7;
void dbg() { cerr << "\\n"; }
template <class F, class... S> void dbg(const F &f, const S &...s) {
  cerr << f << ": ";
  dbg(s...);
}

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  VI p(n), c(n);
  rep(i, n) {
    cin >> p[i];
    p[i]--;
  }
  rep(i, n) cin >> c[i];

  vector<VI> hsu(n, VI(n));
  vector<VL> pt(n, VL(n));
  set<set<int>> st;
  rep(i, n) {
    ll cr = 1, cr2 = 0;
    int z = p[i];
    cr2 = c[z];
    set<int> wst = {z};
    hsu[i][z] = cr;
    pt[i][z] = cr2;
    rep(j, n) {
      z = p[z];
      cr++;
      cr2 += c[z];
      if (hsu[i][z] != 0)
        break;
      hsu[i][z] = cr;
      pt[i][z] = cr2;
      wst.insert(z);
    }
    st.insert(wst);
  }
  VI lp(n);
  VL sk(n + 3);
  for (auto &zz : st) {
    ll sm = 0;
    for (auto z : zz)
      sm += c[z];
    for (auto z : zz) {
      lp[z] = sz(zz);
      sk[z] = sm;
    }
  }

  // for(auto xl: hsu){for(auto x :xl) cerr<<x<<" ";dbg("hsu");}
  // for(auto xl: pt){for(auto x :xl) cerr<<x<<" ";dbg("pt");}
  // for(auto x :sk) cerr<<x<<" "; dbg("sk");
  // for(auto x :lp) cerr<<x<<" "; dbg("lp");
  ll ans = -1LL << 55;
  rep(s, n) rep(e, n) {
    if (hsu[s][e] > k || hsu[s][e] == 0)
      continue;
    ll cnd = pt[s][e];
    ll z = (k - hsu[s][e]) / lp[e];
    cnd += max(0LL, sk[e]) * z;
    ans = max(cnd, ans);
  }
  cout << ans << "\\n";

  return 0;
}