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

// #pragma GCC optimize("O3")
// #pragma GCC optimize("unroll-loops")
// #pragma GCC target("avx2")
//  293206GT
#include <algorithm>
#include <algorithm>
#include <bitset>
#include <chrono>
#include <cmath>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

typedef long long ll;
#define int ll
using namespace std;
#define put(a) cout << (a) << '\\n'
#define sqr(x) (x) * (x)
typedef pair<int, int> pii;
typedef long double ld;
typedef pair<ld, ld> pld;
typedef unsigned long long ull;

typedef vector<int> vi;
typedef vector<pii> vpii;
#define rep(x, y, a) for (int x = (y); x < (int)(a); ++x)
#define all(a) a.begin(), a.end()
#define chkmax(a, b) a = max(a, (b))
#define chkmin(a, b) a = min(a, (b))

namespace IO {

template <class A, class B>
ostream &operator<<(ostream &out, vector<pair<A, B>> a);

template <class A> ostream &operator<<(ostream &out, vector<A> a);

template <class A, class B> ostream &operator<<(ostream &out, pair<A, B> a) {
  out << a.first << " " << a.second;
  return out;
}

template <class A, class B>
ostream &operator<<(ostream &out, vector<pair<A, B>> a) {
  for (pair<A, B> x : a)
    out << x.first << " " << x.second << '\\n';
  return out;
}

template <class A> ostream &operator<<(ostream &out, vector<A> a) {
  for (A x : a)
    out << x << ' ';
  return out;
}

template <class A, class B> istream &operator>>(istream &in, pair<A, B> &a) {
  in >> a.first >> a.second;
  return in;
}

template <class A> istream &operator>>(istream &in, vector<A> &a) {
  for (A &x : a)
    in >> x;
  return in;
}
} // namespace IO
using namespace IO;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

void solve() {
  int n, k;
  cin >> n >> k;
  vi p(n);
  cin >> p;
  vi c(n);
  cin >> c;
  rep(i, 0, n)-- p[i];
  int ans = -1e17;
  rep(st, 0, n) {
    int cur = p[st];
    int an = 0;
    int cnt = 0;
    vi route;
    while (cur != st && cnt < k) {
      route.push_back(c[cur]);
      cur = p[cur];
      ++cnt;
    }
    if (cnt != k)
      route.push_back(c[st]);
    int sum_cir = 0;
    rep(i, 0, route.size()) { sum_cir += route[i]; }
    an = 0;
    rep(i, 0, route.size()) {
      an += route[i];
      chkmax(ans,
             an + max(0ll, (int)(k - i - 1) / (int)route.size() * sum_cir));
    }
  }
  put(ans);
}

signed main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  cout.precision(30);
  //    freopen("input.txt", "r", stdin);
  int q = 1;
  //    cin >> q;

  while (q--)
    solve();
}


Error: TLE

Fixed submission:

// #pragma GCC optimize("O3")
// #pragma GCC optimize("unroll-loops")
// #pragma GCC target("avx2")
//  293206GT
#include <algorithm>
#include <algorithm>
#include <bitset>
#include <chrono>
#include <cmath>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

typedef long long ll;
#define int ll
using namespace std;
#define put(a) cout << (a) << '\\n'
#define sqr(x) (x) * (x)
typedef pair<int, int> pii;
typedef long double ld;
typedef pair<ld, ld> pld;
typedef unsigned long long ull;

typedef vector<int> vi;
typedef vector<pii> vpii;
#define rep(x, y, a) for (int x = (y); x < (int)(a); ++x)
#define all(a) a.begin(), a.end()
#define chkmax(a, b) a = max(a, (b))
#define chkmin(a, b) a = min(a, (b))

namespace IO {

template <class A, class B>
ostream &operator<<(ostream &out, vector<pair<A, B>> a);

template <class A> ostream &operator<<(ostream &out, vector<A> a);

template <class A, class B> ostream &operator<<(ostream &out, pair<A, B> a) {
  out << a.first << " " << a.second;
  return out;
}

template <class A, class B>
ostream &operator<<(ostream &out, vector<pair<A, B>> a) {
  for (pair<A, B> x : a)
    out << x.first << " " << x.second << '\\n';
  return out;
}

template <class A> ostream &operator<<(ostream &out, vector<A> a) {
  for (A x : a)
    out << x << ' ';
  return out;
}

template <class A, class B> istream &operator>>(istream &in, pair<A, B> &a) {
  in >> a.first >> a.second;
  return in;
}

template <class A> istream &operator>>(istream &in, vector<A> &a) {
  for (A &x : a)
    in >> x;
  return in;
}
} // namespace IO
using namespace IO;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

void solve() {
  int n, k;
  cin >> n >> k;
  vi p(n);
  cin >> p;
  vi c(n);
  cin >> c;
  rep(i, 0, n)-- p[i];
  int ans = -1e17;
  rep(st, 0, n) {
    int cur = p[st];
    int an = 0;
    int cnt = 0;
    while (cnt < k) {
      an += c[cur];
      chkmax(ans, an);
      cur = p[cur];
      ++cnt;
    }
  }
  put(ans);
}

signed main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  cout.precision(30);
  //    freopen("input.txt", "r", stdin);
  int q = 1;
  //    cin >> q;

  while (q--)
    solve();
}