alexjercan/bugnet huggingface dataset:
Problem p02585 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

ll dp[5001][5001];
int end_jmp[5001];

ll res = -1e18;

void set_table(vector<int> &pp, vector<int> &cc, int node);

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);

  int n, k;
  cin >> n >> k;

  vector<int> p(n + 1), c(n + 1);

  for (int i = 1; i <= n; ++i)
    cin >> p[i];
  for (int j = 1; j <= n; ++j)
    cin >> c[j];

  for (int s = 1; s <= n; ++s)
    set_table(p, c, s);

  if (k <= 5000) {
    for (int i = 1; i <= n; ++i)
      res = max(res, *max_element(&dp[i][1], &dp[i][1] + k));
  } else {
    for (int i = 1; i <= n; ++i) {
      int cycle = end_jmp[i];
      int dv = k / cycle;
      int mod = k % cycle;

      ll a = -1e18, b = -1e18;
      for (int j = 1; j <= mod; ++j)
        a = max(a, dp[i][j]);
      for (int u = 1; u <= cycle; ++u)
        b = max(b, dp[i][u]);

      res = max(res, a);
      res = max(res, b);
      res = max(res, a + dv * dp[i][cycle]);
      res = max(res, b + (dv - 1) * dp[i][cycle]);
    }
  }

  cout << res;
}

void set_table(vector<int> &pp, vector<int> &cc, int node) {
  int curr = pp[node];
  int jmp = 1;

  int remain = 5000;

  while (remain--) {
    dp[node][jmp] = dp[node][jmp - 1] + cc[curr];
    ++jmp;
    curr = pp[curr];
  }

  int dest = node;
  curr = pp[node];
  jmp = 1;

  while (curr != dest) {
    ++jmp;
    curr = pp[curr];
  }
  end_jmp[node] = jmp;
}

Error: -11

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

ll dp[5001][5001];
int end_jmp[5001];

ll res = -1e18;

void set_table(vector<int> &pp, vector<int> &cc, int node);

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);

  int n, k;
  cin >> n >> k;

  vector<int> p(n + 1), c(n + 1);

  for (int i = 1; i <= n; ++i)
    cin >> p[i];
  for (int j = 1; j <= n; ++j)
    cin >> c[j];

  for (int s = 1; s <= n; ++s)
    set_table(p, c, s);

  if (k <= 5000) {
    for (int i = 1; i <= n; ++i)
      for (int j = 1; j <= k; ++j)
        res = max(res, dp[i][j]);
  } else {
    for (int i = 1; i <= n; ++i) {
      int cycle = end_jmp[i];
      int dv = k / cycle;
      int mod = k % mod;

      ll a = -1e18, b = -1e18;
      for (int j = 1; j <= mod; ++j)
        a = max(a, dp[i][j]);
      for (int u = 1; u <= mod; ++u)
        b = max(b, dp[i][u]);

      res = max(res, a + dv * dp[i][cycle]);
      res = max(res, b + (dv - 1) * dp[i][cycle]);
    }
  }

  cout << res;
}

void set_table(vector<int> &pp, vector<int> &cc, int node) {
  int curr = pp[node];
  int jmp = 1;

  int remain = 5000;

  while (remain--) {
    dp[node][jmp] = dp[node][jmp - 1] + cc[curr];
    ++jmp;
    curr = pp[curr];
  }

  int dest = node;
  curr = pp[node];
  jmp = 1;

  while (curr != dest) {
    ++jmp;
    curr = pp[curr];
  }
  end_jmp[node] = jmp;
}
Problem p02585 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>

long long int N, K;
std::vector<long long int> C, P;

long long int calc_max(std::vector<long long int> loop_vec) {
  long long int sum = 0;
  for (int j = 0; j < loop_vec.size(); j++) {
    sum += loop_vec[j];
  }
  long long int n_loop = loop_vec.size();
  long long int loop = std::max(0ll, (K - n_loop) / n_loop);
  int count = K - loop * n_loop;
  long long int max = loop_vec[0];
  int n_max = 1;
  for (int i = 0; i < n_loop; i++) {
    long long int s = 0;
    for (int j = 0; j < count; j++) {
      s += loop_vec[(i + j) % n_loop];
      if (s > max) {
        max = s;
        n_max = j + 1;
      }
    }
  }
  if (sum > 0) {
    max += sum * loop;
  }
  return max;
}

int main() {
  std::cin >> N >> K;

  C = std::vector<long long int>(N);
  P = std::vector<long long int>(N);
  for (int i = 0; i < N; i++) {
    std::cin >> C[i];
    C[i]--;
  }

  for (int i = 0; i < N; i++) {
    std::cin >> P[i];
  }

  long long int max = -200000000000;
  std::vector<int> flag(N, -1);
  for (int i = 0; i < N; i++) {
    if (flag[i] != -1) {
      continue;
    }
    std::vector<long long int> loop{P[i]};
    flag[i] = 1;
    int next = i;
    while (C[next] != i) {
      next = C[next];
      loop.emplace_back(P[next]);
      flag[next] = 1;
    }
    max = std::max(max, calc_max(loop));
  }
  std::cout << max << std::endl;
}


Error: -6
terminate called after throwing an instance of 'std::exception'
  what():  std::exception


Fixed submission:

#include <bits/stdc++.h>

long long int N, K;
std::vector<long long int> C, P;

long long int calc_max(std::vector<long long int> loop_vec) {
  long long int sum = 0;
  for (int j = 0; j < loop_vec.size(); j++) {
    sum += loop_vec[j];
  }
  int n_loop = loop_vec.size();
  int loop = (K - 1) / n_loop;
  int count = K - loop * n_loop;
  if (count >= n_loop) {
    throw std::exception();
  }
  long long int max = loop_vec[0];
  for (int i = 0; i < n_loop; i++) {
    long long int s = 0;
    for (int j = 0; j < count; j++) {
      s += loop_vec[(i + j) % n_loop];
      max = std::max(max, s);
    }
  }
  if (sum > 0) {
    max += sum * loop;
  }
  return max;
}

int main() {
  std::cin >> N >> K;

  C = std::vector<long long int>(N);
  P = std::vector<long long int>(N);
  for (int i = 0; i < N; i++) {
    std::cin >> C[i];
    C[i]--;
  }

  for (int i = 0; i < N; i++) {
    std::cin >> P[i];
  }

  long long int max = -200000000000;
  std::vector<int> flag(N, -1);
  for (int i = 0; i < N; i++) {
    if (flag[i] != -1) {
      continue;
    }
    std::vector<long long int> loop{P[i]};
    flag[i] = 1;
    int next = i;
    while (C[next] != i) {
      next = C[next];
      loop.emplace_back(P[next]);
      flag[next] = 1;
    }
    max = std::max(max, calc_max(loop));
  }
  std::cout << max << std::endl;
}
Problem p02585 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>

using namespace std;

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);
  int n, k;
  cin >> n >> k;
  int p[n + 1];
  for (int i = 1; i <= n; i++)
    cin >> p[i];
  int c[n + 1];
  for (int i = 1; i <= n; i++)
    cin >> c[i];
  long long m = -2e9;
  for (int i = 1; i <= n; i++) {
    int t = p[i];
    int cyc = 1;
    long long psum[n + 1];
    psum[0] = 0;
    psum[1] = c[i];
    while (t != i) {
      psum[cyc + 1] = psum[cyc] + c[t];
      t = p[t];
      cyc++;
    }
    long long temp = -2e9;
    if (psum[cyc] >= 0 && k >= cyc) {
      temp = psum[cyc] * (k / cyc - 1);
      for (int i = 1; i <= cyc; i++) {
        temp = max(temp, psum[cyc] * (k / cyc - 1) + psum[i]);
      }
      for (int i = 0; i <= k % cyc; i++) {
        temp = max(temp, psum[cyc] * (k / cyc) + psum[i]);
      }
    } else {
      for (int j = 1; j <= min(cyc, k); j++) {
        temp = max(temp, psum[j]);
      }
    }
    m = max(m, temp);
  }
  cout << m;
  return 0;
}

Error: WA

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);
  int n, k;
  cin >> n >> k;
  int p[n + 1];
  for (int i = 1; i <= n; i++)
    cin >> p[i];
  int c[n + 1];
  for (int i = 1; i <= n; i++)
    cin >> c[i];
  long long m = -2e9;
  for (int i = 1; i <= n; i++) {
    int t = p[i];
    int cyc = 1;
    long long psum[n + 1];
    psum[0] = 0;
    psum[1] = c[i];
    while (t != i) {
      psum[cyc + 1] = psum[cyc] + c[t];
      t = p[t];
      cyc++;
    }
    long long temp = -2e9;
    if (psum[cyc] > 0 && k > cyc) {
      temp = psum[cyc] * (k / cyc - 1);
      for (int i = 1; i <= cyc; i++) {
        temp = max(temp, psum[cyc] * (k / cyc - 1) + psum[i]);
      }
      for (int i = 0; i <= k % cyc; i++) {
        temp = max(temp, psum[cyc] * (k / cyc) + psum[i]);
      }
    } else {
      for (int j = 1; j <= max(cyc, k); j++) {
        int num = 0;
        int newj = j;
        newj -= (j / cyc) * k;
        temp = max(temp, psum[newj] + psum[cyc] * (j / cyc));
      }
    }
    m = max(m, temp);
  }
  cout << m;
  return 0;
}
Problem p02585 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
#define inf (long long)(1e18 + 1)
using namespace std;

long long n, k;
vector<long long> c, p;

long long solve();

int main() {
  cin >> n >> k;
  p.resize(n);
  for (auto &x : p) {
    cin >> x;
    --x;
  }
  c.resize(n);
  for (auto &x : c)
    cin >> x;
  cout << solve() << endl;
  return 0;
}

long long solve() {
  long long res = -inf;
  for (int i = 0; i < n; ++i)
    res = max(res, c[p[i]]);
  for (int i = 0; i < n; ++i) {
    vector<bool> ch(n, 0);
    vector<int> v(1, i);
    ch[i] = 1;
    int now = i, cnt = 0;
    long long sum = 0;
    while (cnt < k) {
      now = p[now];
      sum += c[now];
      v.push_back(now);
      res = max(res, sum);
      ++cnt;
      if (ch[now])
        break;
      ch[now] = 1;
    }
    if (cnt == k)
      continue;
    int ls = now;
    long long loops = c[now], lcnt = 1;
    v.pop_back();
    while (v.back() != ls) {
      loops += c[v.back()];
      v.pop_back();
      ++lcnt;
    }
    if (loops >= 0) {
      long long pls = max((k - cnt) / lcnt - 1, 0LL);
      cnt += pls * lcnt;
      sum += loops * pls;
    }
    cnt = max((long long)cnt, k - 3 * n);
    res = max(res, sum);
    while (cnt < k) {
      now = p[now];
      sum += c[now];
      res = max(res, sum);
      ++cnt;
    }
  }
  return res;
}


Error: 0

Fixed submission:

#include <bits/stdc++.h>
#define inf (long long)(1e17)
using namespace std;

long long n, k;
vector<long long> c, p;

long long solve();

int main() {
  cin >> n >> k;
  p.resize(n);
  for (auto &x : p) {
    cin >> x;
    --x;
  }
  c.resize(n);
  for (auto &x : c)
    cin >> x;
  cout << solve() << endl;
  return 0;
}

long long solve() {
  long long res = -inf;
  for (int i = 0; i < n; ++i)
    res = max(res, c[p[i]]);
  for (int i = 0; i < n; ++i) {
    vector<bool> ch(n, 0);
    vector<int> v(1, i);
    ch[i] = 1;
    int now = i, cnt = 0;
    long long sum = 0;
    while (cnt < k) {
      now = p[now];
      sum += c[now];
      v.push_back(now);
      res = max(res, sum);
      ++cnt;
      if (ch[now])
        break;
      ch[now] = 1;
    }
    int ls = now;
    long long loops = c[now], lcnt = 1;
    v.pop_back();
    while (v.back() != ls) {
      loops += c[v.back()];
      v.pop_back();
      ++lcnt;
    }
    if (loops >= 0) {
      long long pls = (k - cnt) / lcnt;
      cnt += pls * lcnt;
      sum += loops * pls;
    }
    res = max(res, sum);
    while (cnt < k && p[now] != ls) {
      now = p[now];
      sum += c[now];
      res = max(res, sum);
      ++cnt;
    }
  }
  return res;
}
Problem p02585 (C++) — buggy submission, verdict: Runtime Error

#ifndef DEBUG
#pragma GCC optimize("O3,no-stack-protector")
#pragma GCC optimize("unroll-loops")

#if __cplusplus < 201703L
#pragma GCC target("avx")
#else
#pragma GCC target("avx2")
#endif

#endif

#define _USE_MATH_DEFINES
#include <bits/stdc++.h>
#define dump(x) cout << x << endl
typedef int64_t Int;
typedef long double Ld;

using namespace std;
using Graph = vector<vector<Int>>;

const Ld pi = M_PI;
const Int MOD = 1000000007;
const Int INF = 1LL << 62;

Int Floor(Int a, Int b) { return (a - (a % b)) / b; }

Int ser(Int n, Int k, const vector<Int> &p, const vector<Int> &c, Int index) {
  Int cycle_size = 0;
  Int cycle_score = 0;

  Int i = index;
  Int ma = -INF;
  while (true) {
    cycle_size++;
    i = p[i];
    cycle_score += c[i];
    if (cycle_size <= k) {
      ma = max(ma, cycle_score);
    }
    if (i == index) {
      break;
    }
  }

  if (cycle_score <= 0) {
    return ma;
  } else {
    Int ret = (k / cycle_size - 1) * cycle_score;
    Int tem = ret;
    i = index;
    ma = -INF;
    Int ro = cycle_size + (k % cycle_size);
    for (Int j = 0; j < ro; j++) {
      i = p[i];
      tem += c[i];
      ret = max(ret, tem);
    }

    return ret;
  }
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);

  Int n, k;
  cin >> n >> k;
  vector<Int> p(n), c(n);
  for (Int i = 0; i < n; i++) {
    cin >> p[i];
    p[i]--;
  }
  for (Int i = 0; i < n; i++) {
    cin >> c[i];
  }

  Int ans = -INF;
  for (Int i = 0; i < n; i++) {
    ans = max(ans, ser(n, k, p, c, i));
  }

  cout << ans << endl;

  return 0;
}

Error: 0

Fixed submission:

#ifndef DEBUG
#pragma GCC optimize("O3,no-stack-protector")
#pragma GCC optimize("unroll-loops")

#if __cplusplus < 201703L
#pragma GCC target("avx")
#else
#pragma GCC target("avx2")
#endif

#endif

#define _USE_MATH_DEFINES
#include <bits/stdc++.h>
#define dump(x) cout << x << endl
typedef int64_t Int;
typedef long double Ld;

using namespace std;
using Graph = vector<vector<Int>>;

const Ld pi = M_PI;
const Int MOD = 1000000007;
const Int INF = 1LL << 61;

Int Floor(Int a, Int b) { return (a - (a % b)) / b; }

Int dfs(Int k, const vector<Int> &p, const vector<Int> &c, Int i) {
  vector<bool> s(p.size(), false);
  vector<Int> psum;
  psum.push_back(0);
  Int nsum = 0;
  Int ma = -((Int)1 << 62);
  while (0 < k) {
    if (s[i]) {
      break;
    }

    s[i] = true;
    i = p[i];
    k--;
    nsum += c[i];
    psum.emplace_back(nsum);
    ma = max(ma, nsum);
  }

  Int n = s.size();
  if (psum.back() < 0) {
    return ma;
  } else {
    return max(ma, psum[k % n]);
  }
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);

  Int n, k;
  cin >> n >> k;
  vector<Int> p(n), c(n);
  for (Int i = 0; i < n; i++) {
    cin >> p[i];
    p[i]--;
  }
  for (Int i = 0; i < n; i++) {
    cin >> c[i];
  }

  Int ans = 0;
  for (Int i = 0; i < n; i++) {
    Int tem = 0;

    tem = dfs(k, p, c, i);
    ans = max(ans, tem);
  }

  if (ans == 0) {
    sort(c.begin(), c.end(), greater<>());
    ans = c.front();
  }

  cout << ans << endl;

  return 0;
}
Problem p02585 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <algorithm>
#include <iostream>
#include <string>
#include <vector>

using namespace std;
using ll = long long;
ll INF = 1LL << 60;

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> P(N), C(N);
  for (int i = 0; i < N; i++) {
    cin >> P[i];
    P[i]--;
  }
  for (int i = 0; i < N; i++)
    cin >> C[i];

  ll ans = -INF;

  for (int i = 0; i < N; i++) {
    int v = i;
    ll cycle_sum = 0;
    int cycle_cnt = 0;

    while (true) {
      cycle_cnt++;
      cycle_sum += C[v];
      v = P[v];
      if (v == i)
        break;
    }

    ll remainder = 0;
    int cnt = 0;

    while (true) {
      cnt++;
      remainder += C[v];

      if (cnt > K)
        break;

      int roop = (K - cnt) / cycle_cnt;
      ll score = remainder + max(0LL, cycle_sum) * roop;
      ans = max(ans, score);

      v = P[v];
      if (v == i)
        break;
    }
  }

  cout << ans << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <string>
#include <vector>

using namespace std;
using ll = long long;
const ll INF = 1LL << 60;

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> P(N), C(N);
  for (int i = 0; i < N; i++) {
    cin >> P[i];
    P[i]--;
  }
  for (int i = 0; i < N; i++)
    cin >> C[i];

  ll ans = -INF;

  for (int i = 0; i < N; i++) {
    int v = i;
    ll cycle_sum = 0;
    int cycle_cnt = 0;

    while (true) {
      cycle_cnt++;
      cycle_sum += C[v];
      v = P[i];
      if (v == i)
        break;
    }

    ll remainder = 0;
    int cnt = 0;

    while (true) {
      cnt++;
      remainder += C[v];

      if (cnt > K)
        break;

      int roop = (K - cnt) / cycle_cnt;
      ll score = remainder + max(0LL, cycle_sum) * roop;
      ans = max(ans, score);

      v = P[v];
      if (v == i)
        break;
    }
  }

  cout << ans << endl;

  return 0;
}
Problem p02585 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <iostream>
#include <vector>
using namespace std;
int main() {
  long n, k;
  cin >> n >> k;
  vector<long> p(n), c(n);
  for (long i = 0; i < n; i++) {
    cin >> p[i];
    p[i]--;
  }
  for (long i = 0; i < n; i++) {
    cin >> c[i];
  }
  long ms = c[p[0]];
  for (long i = 0; i < n; i++) {
    long score = 0;
    long tk = k;
    long x = i;
    bool loop = false;
    while (tk > 0) {
      tk--;
      x = p[x];
      score += c[x];
      ms = max(ms, score);
      if (x == i && !loop) {
        loop = true;
        long t = tk / (k - tk) - 1;
        if (t > 0) {
          score += score * t;
          tk -= (k - tk) * t;
        }
        ms = max(ms, score);
      }
    }
  }
  cout << ms << endl;
}


Error: TLE

Fixed submission:

#include <iostream>
#include <vector>
using namespace std;
int main() {
  long n, k;
  cin >> n >> k;
  vector<long> p(n), c(n);
  for (long i = 0; i < n; i++) {
    cin >> p[i];
    p[i]--;
  }
  for (long i = 0; i < n; i++) {
    cin >> c[i];
  }
  long ms = c[p[0]];
  bool loop = false;
  for (long i = 0; i < n; i++) {
    long score = 0;
    long tk = k;
    long x = i;
    while (tk > 0) {
      tk--;
      x = p[x];
      score += c[x];
      ms = max(ms, score);
      if (x == i && !loop) {
        loop = true;
        long t = tk / (k - tk) - 1;
        if (t > 0) {
          score += score * t;
          tk -= (k - tk) * t;
        }
        ms = max(ms, score);
      }
    }
  }
  cout << ms << endl;
}
Problem p02585 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
using namespace std;

using ll = long long;
using pii = pair<int, int>;
template <class T> using V = vector<T>;
template <class T> using VV = V<V<T>>;

#define pb push_back
#define eb emplace_back
#define mp make_pair
#define fi first
#define se second
#define rep(i, n) rep2(i, 0, n)
#define rep2(i, m, n) for (int i = m; i < (n); i++)
#define per(i, b) per2(i, 0, b)
#define per2(i, a, b) for (int i = int(b) - 1; i >= int(a); i--)
#define ALL(c) (c).begin(), (c).end()
#define SZ(x) ((int)(x).size())

constexpr ll TEN(int n) { return (n == 0) ? 1 : 10 * TEN(n - 1); }

template <class T, class U> void chmin(T &t, const U &u) {
  if (t > u)
    t = u;
}
template <class T, class U> void chmax(T &t, const U &u) {
  if (t < u)
    t = u;
}

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

template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "{";
  rep(i, v.size()) {
    if (i)
      os << ",";
    os << v[i];
  }
  os << "}";
  return os;
}

#ifdef LOCAL
void debug_out() { cerr << endl; }
template <typename Head, typename... Tail> void debug_out(Head H, Tail... T) {
  cerr << " " << H;
  debug_out(T...);
}
#define debug(...)                                                             \\
  cerr << __LINE__ << " [" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#define dump(x) cerr << __LINE__ << " " << #x << " = " << (x) << endl
#else
#define debug(...) (void(0))
#define dump(x) (void(0))
#endif

constexpr ll INF = TEN(18);

int main() {
  int N, K;
  cin >> N >> K;
  V<int> P(N), C(N);
  rep(i, N) cin >> P[i], --P[i];
  rep(i, N) cin >> C[i];

  ll ans = -INF;

  rep(i, N) {
    int v = i;
    ll cycle_sum = 0;
    int cycle_cnt = 0;

    while (true) {
      cycle_cnt++;
      cycle_sum += C[v];
      v = P[v];
      if (v == i)
        break;
    }

    ll path = 0;
    int cnt = 0;

    while (true) {
      cnt++;
      path += C[v];

      if (cnt > K)
        break;

      int num = (K - cnt) / cycle_cnt;
      ll score = path + max(0ll, cycle_sum) * num;
      chmax(ans, score);

      v = P[v];
      if (v == i)
        break;
    }
  }

  cout << ans << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define ll long long
using namespace std;

int main() {
  int n, k;
  cin >> n >> k;
  vector<int> a(n + 1);
  vector<ll> b(n + 1);
  for (int i = 1; i <= n; i++)
    cin >> a[i];
  for (int i = 1; i <= n; i++)
    cin >> b[i];
  ll ans = b[1];
  for (int i = 1; i <= n; i++) {
    ll c = 0;
    int in = i;
    for (int j = 0; j < k; j++) {
      in = a[in];
      c += b[in];
      ans = max(ans, c);
    }
  }
  cout << ans << '\\n';
  return 0;
}
Problem p02585 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <algorithm>
#include <stdio.h>
using namespace std;

typedef long long int ll;
constexpr int kN = int(5E3 + 10), LOGN = 31;
constexpr ll kInf = ll(3E18 + 10);

int p[LOGN][kN];
ll c[kN], val[LOGN][kN], ans[kN];

int main() {
  int n, k, npos, lim;
  ll a = -kInf, tmp;
  scanf("%d%d", &n, &k);
  for (int i = 1; i <= n; i++)
    scanf("%d", &p[0][i]);
  for (int i = 1; i <= n; i++)
    scanf("%lld", &c[i]);
  for (int i = 1; i <= n; i++)
    val[0][i] = c[p[0][i]];
  for (int i = 1; i < LOGN; i++)
    for (int j = 1; j <= n; j++) {
      p[i][j] = p[i - 1][p[i - 1][j]];
      val[i][j] = val[i - 1][j] + val[i - 1][p[i - 1][j]];
    }

  for (int i = 1; i <= n; i++) {
    npos = i;
    for (int j = LOGN - 1; j >= 0; j--)
      if (k & (1 << j)) {
        ans[i] += val[j][npos];
        npos = p[j][npos];
      }
  }

  lim = min(n, k);
  for (int kk = 1; kk <= lim; kk++) {
    for (int i = 1; i <= n; i++) {
      npos = i;
      tmp = 0;
      for (int j = LOGN - 1; j >= 0; j--)
        if (kk & (1 << j)) {
          tmp += val[j][npos];
          npos = p[j][npos];
        }
      ans[i] = max(ans[i], tmp);
    }
  }
  lim = max(1, k - n);
  for (int kk = lim; kk <= k; kk++) {
    for (int i = 1; i <= n; i++) {
      npos = i;
      tmp = 0;
      for (int j = LOGN - 1; j >= 0; j--)
        if (kk & (1 << j)) {
          tmp += val[j][npos];
          npos = p[j][npos];
        }
      ans[i] = max(ans[i], tmp);
    }
  }

  for (int i = 1; i <= n; i++)
    a = max(a, ans[i]);
  printf("%lld\\n", a);
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <stdio.h>
using namespace std;

typedef long long int ll;
constexpr int kN = int(5E3 + 10), LOGN = 31;
constexpr ll kInf = ll(3E18 + 10);

int p[LOGN][kN];
ll c[kN], val[LOGN][kN], ans[kN];

int main() {
  int n, k, npos, lim;
  ll a = -kInf, tmp;
  scanf("%d%d", &n, &k);
  for (int i = 1; i <= n; i++)
    scanf("%d", &p[0][i]);
  for (int i = 1; i <= n; i++)
    scanf("%lld", &c[i]);
  for (int i = 1; i <= n; i++)
    val[0][i] = c[p[0][i]];
  for (int i = 1; i < LOGN; i++)
    for (int j = 1; j <= n; j++) {
      p[i][j] = p[i - 1][p[i - 1][j]];
      val[i][j] = val[i - 1][j] + val[i - 1][p[i - 1][j]];
    }

  for (int i = 1; i <= n; i++) {
    npos = i;
    for (int j = LOGN - 1; j >= 0; j--)
      if (k & (1 << j)) {
        ans[i] += val[j][npos];
        npos = p[j][npos];
      }
  }

  lim = min(n << 1, k);
  for (int kk = 1; kk <= lim; kk++) {
    for (int i = 1; i <= n; i++) {
      npos = i;
      tmp = 0;
      for (int j = LOGN - 1; j >= 0; j--)
        if (kk & (1 << j)) {
          tmp += val[j][npos];
          npos = p[j][npos];
        }
      ans[i] = max(ans[i], tmp);
    }
  }
  lim = max(1, min(k - (n << 1), k));
  for (int kk = lim; kk <= k; kk++) {
    for (int i = 1; i <= n; i++) {
      npos = i;
      tmp = 0;
      for (int j = LOGN - 1; j >= 0; j--)
        if (kk & (1 << j)) {
          tmp += val[j][npos];
          npos = p[j][npos];
        }
      ans[i] = max(ans[i], tmp);
    }
  }

  for (int i = 1; i <= n; i++)
    a = max(a, ans[i]);
  printf("%lld\\n", a);
}