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

#include <algorithm>
#include <assert.h>
#include <bitset>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <stdio.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

#define pb push_back
#define fi first
#define se second
#define all(x) (x).begin(), (x).end()
#define mn(x) *min_element((x).begin(), (x).end())
#define mx(x) *max_element((x).begin(), (x).end())
#define eb emplace_back
#define el '\\n'
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef pair<ll, ll> pll;
typedef pair<int, int> pii;
typedef vector<ll> vll;
typedef vector<int> vi;
typedef vector<pii> vpi;
typedef vector<pll> vpll;
typedef vector<vi> vvi;
typedef vector<vll> vvll;
typedef vector<bool> vb;
typedef vector<string> vs;

const ld pi = acos(-1);
const ll inf = (ll)1e18;
const int mod = 1000000007;
// const int mod=998244353;

int n, m, k;
const int N = 1e6 + 10;

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  int q, t, u, v, h;
  ll x, a, b, c, d, y, w, z, k;
  cin >> n >> k;
  vi pp(n + 1), vv(n + 1);
  for (int i = 1; i <= n; i++) {
    cin >> pp[i];
  }
  for (int i = 1; i <= n; i++) {
    cin >> vv[i];
  }
  vvll dp(n + 1);

  auto f = [&](int id) {
    int n = (int)dp[id].size();
    for (int i = 1; i < n; i++) {
      dp[id][i] += dp[id][i - 1];
    }
    ll ans;
    if (k > n) {
      m = k % n;
      ans = 0;
      for (int i = 0; i < m; i++) {
        ans = max(ans, dp[id][i]);
      }
      ans = max({mx(dp[id]), (k / n) * dp[id][n - 1] + ans,
                 (k / n - 1) * dp[id][n - 1] + mx(dp[id])});
    } else {
      ans = dp[id][0];
      for (int i = 0; i < k; i++) {
        ans = max(ans, dp[id][i]);
      }
    }
    return ans;
  };
  for (int i = 1; i <= n; i++) {
    int st = i;
    vi vis(n + 1, -1);
    while (true) {
      if (~vis[st]) {
        break;
      } else {
        vis[st] = 1;
        st = pp[st];
        dp[i].pb(vv[st]);
      }
    }
  }
  a = -1e18;
  for (int i = 1; i <= n; i++) {
    a = max(a, f(i));
  }
  cout << a << el;

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <assert.h>
#include <bitset>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <stdio.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

#define pb push_back
#define fi first
#define se second
#define all(x) (x).begin(), (x).end()
#define mn(x) *min_element((x).begin(), (x).end())
#define mx(x) *max_element((x).begin(), (x).end())
#define eb emplace_back
#define el '\\n'
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef pair<ll, ll> pll;
typedef pair<int, int> pii;
typedef vector<ll> vll;
typedef vector<int> vi;
typedef vector<pii> vpi;
typedef vector<pll> vpll;
typedef vector<vi> vvi;
typedef vector<vll> vvll;
typedef vector<bool> vb;
typedef vector<string> vs;

const ld pi = acos(-1);
const ll inf = (ll)1e18;
const int mod = 1000000007;
// const int mod=998244353;

int n, m, k;
const int N = 1e6 + 10;

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  int q, t, u, v, h;
  ll x, a, b, c, d, y, w, z, k;
  cin >> n >> k;
  vi pp(n + 1), vv(n + 1);
  for (int i = 1; i <= n; i++) {
    cin >> pp[i];
  }
  for (int i = 1; i <= n; i++) {
    cin >> vv[i];
  }
  a = -1e18;
  for (int i = 1; i <= n; i++) {
    int st = i;
    b = 0;
    for (int i = 1; i <= k; i++) {
      st = pp[st];
      if (i == 1)
        c = vv[st];
      b += vv[st];
      c = max(c, b);
    }
    a = max(a, c);
  }
  cout << a << el;

  return 0;
}