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

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
typedef pair<ll, ll> P;
typedef pair<int, int> Pi;
#define rep(i, n) for (ll i = 0; i < n; i++)
#define FOR(i, a, b) for (ll i = a; i < b; i++)
#define fi first
#define se second
#define endl "\\n"

template <typename T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> ostream &operator<<(ostream &s, const complex<T> &d) {
  return s << "(" << d.real() << ", " << d.imag() << ")";
}
template <typename T1, typename T2>
ostream &operator<<(ostream &s, const pair<T1, T2> &d) {
  return s << "(" << d.first << ", " << d.second << ")";
}
template <typename T> ostream &operator<<(ostream &s, const vector<T> &d) {
  int len = d.size();
  rep(i, len) {
    s << d[i];
    if (i < len - 1)
      s << " ";
  }
  return s;
}
template <typename T>
ostream &operator<<(ostream &s, const vector<vector<T>> &d) {
  int len = d.size();
  rep(i, len) { s << d[i] << endl; }
  return s;
}
template <typename T> ostream &operator<<(ostream &s, const set<T> &v) {
  s << "{ ";
  for (auto itr = v.begin(); itr != v.end(); ++itr) {
    if (itr != v.begin()) {
      s << ", ";
    }
    s << (*itr);
  }
  s << " }";
  return s;
}
template <typename T> ostream &operator<<(ostream &s, const multiset<T> &v) {
  s << "{ ";
  for (auto itr = v.begin(); itr != v.end(); ++itr) {
    if (itr != v.begin()) {
      s << ", ";
    }
    s << (*itr);
  }
  s << " }";
  return s;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &s, const map<T1, T2> &m) {
  s << "{" << endl;
  for (auto itr = m.begin(); itr != m.end(); ++itr) {
    s << " " << (*itr).first << " : " << (*itr).second << endl;
  }
  s << "}" << endl;
  return s;
}

const ll mod = 1'000'000'007;
const ll inf = 1'000'000'000'000'000'00;
const int INF = 1'000'000'000;
const double EPS = 1e-10;
const double PI = acos(-1);

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  ll n, k;
  cin >> n >> k;
  vector<ll> p(n), c(n);
  rep(i, n) {
    cin >> p[i];
    p[i]--;
  }
  bool flag = true;
  rep(i, n) {
    cin >> c[i];
    if (c[i] > 0)
      flag = false;
  }
  if (flag) { // 全部負だったら
    cout << *max_element(begin(c), end(c)) << endl;
    return 0;
  }
  // 答えは絶対正になる
  ll ans = 0;
  rep(st, n) {
    queue<ll> que;
    que.push(st);
    vector<ll> dp(n, -inf);
    dp[st] = 0;
    ll cnt = 0;
    while (!que.empty()) {
      auto v = que.front();
      que.pop();
      if (cnt == k)
        break;
      ll to = p[v];
      if (dp[to] == -inf) {
        que.push(to);
        cnt++;
        dp[to] = dp[v] + c[to];
      } else {
        ll roop = dp[v] + c[to] - dp[to];
        if (roop < 0)
          break;
        dp[to] += roop;
        cnt++;
        ll num = 0;
        for (ll i = to;; i = p[i]) {
          num++;
          if (i == v)
            break;
        }
        ll ret = k - cnt;
        if (ret >= num) {
          dp[to] += roop * (ret / num - 1);
          ret %= num;
          ret += num;
        }
        ll now = to;
        while (ret--) {
          dp[p[now]] = dp[now] + c[p[now]];
          now = p[now];
          chmax(ans, dp[now]);
        }
        break;
      }
    }
    rep(i, n) { chmax(ans, dp[i]); }
  }
  cout << ans << endl;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
typedef pair<ll, ll> P;
typedef pair<int, int> Pi;
#define rep(i, n) for (ll i = 0; i < n; i++)
#define FOR(i, a, b) for (ll i = a; i < b; i++)
#define fi first
#define se second
#define endl "\\n"

template <typename T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> ostream &operator<<(ostream &s, const complex<T> &d) {
  return s << "(" << d.real() << ", " << d.imag() << ")";
}
template <typename T1, typename T2>
ostream &operator<<(ostream &s, const pair<T1, T2> &d) {
  return s << "(" << d.first << ", " << d.second << ")";
}
template <typename T> ostream &operator<<(ostream &s, const vector<T> &d) {
  int len = d.size();
  rep(i, len) {
    s << d[i];
    if (i < len - 1)
      s << " ";
  }
  return s;
}
template <typename T>
ostream &operator<<(ostream &s, const vector<vector<T>> &d) {
  int len = d.size();
  rep(i, len) { s << d[i] << endl; }
  return s;
}
template <typename T> ostream &operator<<(ostream &s, const set<T> &v) {
  s << "{ ";
  for (auto itr = v.begin(); itr != v.end(); ++itr) {
    if (itr != v.begin()) {
      s << ", ";
    }
    s << (*itr);
  }
  s << " }";
  return s;
}
template <typename T> ostream &operator<<(ostream &s, const multiset<T> &v) {
  s << "{ ";
  for (auto itr = v.begin(); itr != v.end(); ++itr) {
    if (itr != v.begin()) {
      s << ", ";
    }
    s << (*itr);
  }
  s << " }";
  return s;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &s, const map<T1, T2> &m) {
  s << "{" << endl;
  for (auto itr = m.begin(); itr != m.end(); ++itr) {
    s << " " << (*itr).first << " : " << (*itr).second << endl;
  }
  s << "}" << endl;
  return s;
}

const ll mod = 1'000'000'007;
const ll inf = 1'000'000'000'000'000'00;
const int INF = 1'000'000'000;
const double EPS = 1e-10;
const double PI = acos(-1);

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  ll n, k;
  cin >> n >> k;
  vector<ll> p(n), c(n);
  rep(i, n) {
    cin >> p[i];
    p[i]--;
  }
  bool flag = true;
  rep(i, n) {
    cin >> c[i];
    if (c[i] > 0)
      flag = false;
  }
  if (flag) { // 全部負だったら
    cout << *max_element(begin(c), end(c)) << endl;
    return 0;
  }
  // 答えは絶対正になる
  ll ans = 0;
  rep(st, n) {
    queue<ll> que;
    que.push(st);
    vector<ll> dp(n, -inf);
    dp[st] = 0;
    ll cnt = 0;
    while (!que.empty()) {
      auto v = que.front();
      que.pop();
      if (cnt == k)
        break;
      ll to = p[v];
      if (dp[to] == -inf) {
        que.push(to);
        cnt++;
        dp[to] = dp[v] + c[to];
      } else {
        ll roop = dp[v] + c[to] - dp[to];
        if (roop < 0)
          break;
        flag = true;
        dp[to] += roop;
        cnt++;
        ll num = 0;
        for (ll i = to;; i = p[i]) {
          num++;
          if (i == v)
            break;
        }
        ll ret = k - cnt;
        dp[to] += roop * (ret / num);
        ret %= num;
        ll now = to;
        while (ret--) {
          dp[p[now]] = dp[now] + c[p[now]];
          now = p[now];
        }
        for (ll i = now;; i = p[now]) {
          if (p[i] == to)
            break;
          dp[p[i]] = dp[i] + c[p[i]] - roop;
        }
        break;
      }
    }
    rep(i, n) { chmax(ans, dp[i]); }
  }
  cout << ans << endl;
}