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

#include <algorithm>
#include <cassert>
#include <chrono>
#include <climits>
#include <cmath>
#include <complex>
#include <cstring>
#include <fstream>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <regex>
#include <set>
#include <stack>
#include <vector>
using namespace std;

#pragma region

#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define repr(i, s, n) for (int i = (s); i < (int)(n); i++)
#define revrep(i, n) for (int i = (n)-1; i >= 0; i--)
#define revrepr(i, s, n) for (int i = (n)-1; i >= s; i--)
#define debug(x) cerr << #x << ": " << x << "\\n"
#define popcnt(x) __builtin_popcount(x)

const double PI = acos(-1.0);

using ll = long long;
using P = pair<int, int>;

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

template <class T> istream &operator>>(istream &is, vector<T> &v) {
  for (int i = 0; i < (int)v.size(); i++)
    cin >> v.at(i);
  return is;
}

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

template <class T>
void print(const vector<T> &v, const string &delimiter = " ") {
  rep(i, v.size()) cout << (0 < i ? delimiter : "") << v.at(i);
  cout << endl;
}

template <class T>
void print(const vector<vector<T>> &vv, const string &delimiter) {
  for (const auto &v : vv)
    print(v, delimiter);
}

#pragma endregion

int main() {
  ll n, k;
  cin >> n >> k;
  vector<int> p(n);
  vector<ll> c(n);
  cin >> p >> c;
  rep(i, n) p[i]--;

  ll ans = -1e18;
  rep(s, n) {
    vector<bool> done(n);
    ll len = 0;
    ll pos = s;
    ll sum = 0;
    while (1) {
      if (done[pos])
        break;
      if (len == k)
        break;

      done[pos] = true;
      len++;
      sum += c[p[pos]];
      chmax(ans, sum);
      pos = p[pos];
    }

    if (len == k)
      continue;
    if (sum < 0)
      continue;

    ll rem = k % len + len;
    chmax(ans, sum);
    pos = s;
    ll score = (k / len - 1) * sum;
    chmax(ans, score);
    rep(i, rem) {
      score += c[p[pos]];
      pos = p[pos];
      chmax(ans, score);
    }
  }
  cout << ans << endl;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <chrono>
#include <climits>
#include <cmath>
#include <complex>
#include <cstring>
#include <fstream>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <regex>
#include <set>
#include <stack>
#include <vector>
using namespace std;

#pragma region

#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define repr(i, s, n) for (int i = (s); i < (int)(n); i++)
#define revrep(i, n) for (int i = (n)-1; i >= 0; i--)
#define revrepr(i, s, n) for (int i = (n)-1; i >= s; i--)
#define debug(x) cerr << #x << ": " << x << "\\n"
#define popcnt(x) __builtin_popcount(x)

const double PI = acos(-1.0);

using ll = long long;
using P = pair<int, int>;

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

template <class T> istream &operator>>(istream &is, vector<T> &v) {
  for (int i = 0; i < (int)v.size(); i++)
    cin >> v.at(i);
  return is;
}

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

template <class T>
void print(const vector<T> &v, const string &delimiter = " ") {
  rep(i, v.size()) cout << (0 < i ? delimiter : "") << v.at(i);
  cout << endl;
}

template <class T>
void print(const vector<vector<T>> &vv, const string &delimiter) {
  for (const auto &v : vv)
    print(v, delimiter);
}

#pragma endregion

int main() {
  int n, k;
  cin >> n >> k;
  vector<int> p(n);
  vector<ll> c(n);
  cin >> p >> c;
  rep(i, n) p[i]--;

  ll ans = -1e18;
  rep(s, n) {
    vector<bool> done(n);
    int len = 0;
    int pos = s;
    ll sum = 0;
    while (1) {
      if (done[pos])
        break;
      if (len == k)
        break;

      len++;
      sum += c[p[pos]];
      chmax(ans, sum);
      pos = p[pos];
    }

    int rem = k % len;
    chmax(ans, sum);
    pos = s;
    ll score = k / len * sum;
    rep(i, rem) {
      score += c[p[pos]];
      pos = p[pos];
      chmax(ans, score);
    }
  }
  cout << ans << endl;
}