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

// #define _GLIBCXX_DEBUG
// #pragma GCC optimize ("trapv") // kills the program on integer overﬂows (but
// is really slow).
#include <bits/stdc++.h>
using namespace std;
using namespace std::chrono;

#define ll long long
#define ull unsigned long long
#define FIO                                                                    \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL);
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define ff first
#define ss second
#define vt vector
#define vll vt<ll>
#define pll pair<ll, ll>
#define vpll vt<pll>
#define vvll vt<vll>
#define all(v) v.begin(), v.end()
#define FOR(i, n) for (ll i = 0; i < n; i++)
#define ffo(i, a, b) for (ll i = a; i <= b; i++)
#define rfo(i, a, b) for (ll i = a; i >= b; i--)
#define space cout << "\\n\\n";
#define endl "\\n"
#define pqmx priority_queue<ll>
#define pqmn priority_queue<ll, vll, greater<ll>>
#define fps(x, y) fixed << setprecision(y) << x
#define merg(a, b, c)                                                          \\
  set_union(a.begin(), a.end(), b.begin(), b.end(), inserter(c, c.begin()))
#define set_ar(arr, v) memset(arr, v, sizeof(arr))
#define go_t                                                                   \\
  int testcases;                                                               \\
  cin >> testcases;                                                            \\
  ffo(caseno, 1, testcases)

#define ctime auto start = high_resolution_clock::now()
#define etime auto stop = high_resolution_clock::now()
#define ptime                                                                  \\
  auto z1z = duration_cast<microseconds>(stop - start);                        \\
  cout << "Time elapsed : " << z1z.count() << " microseconds\\n"

#define dbg1(x) cout << #x << ": " << x << '\\n';
#define dbg2(x, y) cout << #x << ": " << x << " | " << #y << ": " << y << '\\n';
#define dbg3(x, y, z)                                                          \\
  cout << #x << ": " << x << " | " << #y << ": " << y << " | " << #z << ": "   \\
       << z << '\\n';
#define dbg4(a, b, c, d)                                                       \\
  cout << #a << ": " << a << " | " << #b << ": " << b << " | " << #c << ": "   \\
       << c << " | " << #d << ": " << d << '\\n';
#define dbg5(a, b, c, d, e)                                                    \\
  cout << #a << ": " << a << " | " << #b << ": " << b << " | " << #c << ": "   \\
       << c << " | " << #d << ": " << d << " | " << #e << ": " << e << '\\n';

const ll mod = 1e9 + 7;
const ll N = 5e3 + 6;
const ll maxN = 5e3 + 5;
const ll MAX_SIZE = 2e6 + 6;
const ll INF = 0x3f3f3f3f3f3f3f3fll;
const double PI = 3.14159265359;

ll powerM(ll x, ll y, ll M = mod) { // default argument
  ll v = 1;
  x = x % M;
  while (y > 0) {
    if (y & 1)
      v = (v * x) % M;
    y = y >> 1;
    x = (x * x) % M;
  }
  return v;
}

ll power(ll x, ll y) {
  ll v = 1;
  while (y > 0) {
    if (y & 1)
      v = v * x;
    y = y >> 1;
    x = x * x;
  }
  return v;
}

vt<bool> vis(maxN, 0);
vvll adj(maxN);
vll p(maxN), c(maxN);

void dfs(ll node, ll par, ll &a1, vll &fthis) {
  vis[node] = 1;
  a1 += c[node];
  fthis.pb(a1);
  for (auto &u : adj[node]) {
    if (vis[u])
      continue;
    // a1 += c[u];
    dfs(u, node, a1, fthis);
  }
  vis[node] = 0;
}

int main() {
#ifndef ONLINE_JUDGE
  freopen("in4.txt", "r", stdin);
  freopen("out3.txt", "w", stdout);
#endif
  FIO ll n, k;
  cin >> n >> k;
  FOR(i, n) {
    cin >> p[i];
    adj[i].pb(p[i] - 1);
  }
  FOR(i, n) cin >> c[i];
  // FOR(i, n) {
  // 	for (auto& it : adj[i])
  // 		cout << it << " "; cout << endl;
  // }
  ll ans = -1e16;
  for (int i = 0; i < n; ++i) {
    // if (vis[i])
    // 	continue;
    ll temp = 0;
    ll a1 = 0;
    vll fthis;
    dfs(i, -1, a1, fthis);
    //
    ll sz = fthis.size();
    ll cycles = k / sz;
    if (cycles == 0) {
      for (int j = 0; j < min(k, sz); j++) {
        ans = max(ans, fthis[j]);
      }
    } else {
      for (int j = 0; j < sz; j++) {
        ans = max(ans, fthis[j]);
      }
      temp = (cycles - 1) * fthis.back();
      ll rem = k - ((cycles - 1) * sz);
      if (rem >= sz) {
        ll t2 = temp;
        for (int j = 0; j < sz; ++j) {
          t2 = temp + fthis[j];
          ans = max(t2, ans);
        }
        temp += fthis.back();
        ans = max(ans, temp);
        rem -= sz;
        t2 = temp;
        for (int j = 0; j < rem; ++j) {
          t2 = temp + fthis[j];
          ans = max(t2, ans);
        }
      } else {
        ll t2 = temp;
        for (int j = 0; j < rem; ++j) {
          t2 = temp + fthis[j];
          ans = max(t2, ans);
        }
      }
    }
  }
  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

// #define _GLIBCXX_DEBUG
// #pragma GCC optimize ("trapv") // kills the program on integer overﬂows (but
// is really slow).
#include <bits/stdc++.h>
using namespace std;
using namespace std::chrono;

#define ll long long
#define ull unsigned long long
#define FIO                                                                    \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL);
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define ff first
#define ss second
#define vt vector
#define vll vt<ll>
#define pll pair<ll, ll>
#define vpll vt<pll>
#define vvll vt<vll>
#define all(v) v.begin(), v.end()
#define FOR(i, n) for (ll i = 0; i < n; i++)
#define ffo(i, a, b) for (ll i = a; i <= b; i++)
#define rfo(i, a, b) for (ll i = a; i >= b; i--)
#define space cout << "\\n\\n";
#define endl "\\n"
#define pqmx priority_queue<ll>
#define pqmn priority_queue<ll, vll, greater<ll>>
#define fps(x, y) fixed << setprecision(y) << x
#define merg(a, b, c)                                                          \\
  set_union(a.begin(), a.end(), b.begin(), b.end(), inserter(c, c.begin()))
#define set_ar(arr, v) memset(arr, v, sizeof(arr))
#define go_t                                                                   \\
  int testcases;                                                               \\
  cin >> testcases;                                                            \\
  ffo(caseno, 1, testcases)

#define ctime auto start = high_resolution_clock::now()
#define etime auto stop = high_resolution_clock::now()
#define ptime                                                                  \\
  auto z1z = duration_cast<microseconds>(stop - start);                        \\
  cout << "Time elapsed : " << z1z.count() << " microseconds\\n"

#define dbg1(x) cout << #x << ": " << x << '\\n';
#define dbg2(x, y) cout << #x << ": " << x << " | " << #y << ": " << y << '\\n';
#define dbg3(x, y, z)                                                          \\
  cout << #x << ": " << x << " | " << #y << ": " << y << " | " << #z << ": "   \\
       << z << '\\n';
#define dbg4(a, b, c, d)                                                       \\
  cout << #a << ": " << a << " | " << #b << ": " << b << " | " << #c << ": "   \\
       << c << " | " << #d << ": " << d << '\\n';
#define dbg5(a, b, c, d, e)                                                    \\
  cout << #a << ": " << a << " | " << #b << ": " << b << " | " << #c << ": "   \\
       << c << " | " << #d << ": " << d << " | " << #e << ": " << e << '\\n';

const ll mod = 1e9 + 7;
const ll N = 3e3 + 6;
const ll maxN = 1e4 + 5;
const ll MAX_SIZE = 2e6 + 6;
const ll INF = 0x3f3f3f3f3f3f3f3fll;
const double PI = 3.14159265359;

ll powerM(ll x, ll y, ll M = mod) { // default argument
  ll v = 1;
  x = x % M;
  while (y > 0) {
    if (y & 1)
      v = (v * x) % M;
    y = y >> 1;
    x = (x * x) % M;
  }
  return v;
}

ll power(ll x, ll y) {
  ll v = 1;
  while (y > 0) {
    if (y & 1)
      v = v * x;
    y = y >> 1;
    x = x * x;
  }
  return v;
}

int main() {
#ifndef ONLINE_JUDGE
  freopen("in4.txt", "r", stdin);
  freopen("out3.txt", "w", stdout);
#endif
  FIO ll n, k;
  cin >> n >> k;
  vll p(n), c(n);
  FOR(i, n) {
    cin >> p[i];
    --p[i];
  }
  FOR(i, n) cin >> c[i];
  ll ans = -1e16;
  for (int i = 0; i < n; ++i) {
    ll tot = 0;
    ll next = i;
    for (int j = 0; j < k; ++j) {
      next = p[next];
      tot += c[next];
      ans = max(ans, tot);
    }
  }
  cout << ans;
  return 0;
}