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

#include <algorithm>
#include <bitset>
#include <cmath>
#include <iomanip> //double精度 setprecision
#include <iostream>
#include <map>
#include <numeric> //lcm
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>

// #include <atcoder/all>

using namespace std;
// using namespace atcoder;

#define rep(i, n) for (int i = 0; i < (n); ++i)
#define rrep(i, n) for (int i = n - 1; i >= 0; --i)
#define rep1(i, n) for (int i = 1; i <= (n); ++i)
#define rrep1(i, n) for (int i = (n); i >= 1; --i)
#define REP(i, n, m) for (int i = (n); i < (m); ++i)
#define all(vec) (vec).begin(), (vec).end()

#define debug(vec)                                                             \\
  for (auto v : vec)                                                           \\
    cerr << v << " ";                                                          \\
  cerr << endl;
#define debug2D(vec2D, w)                                                      \\
  for (auto vec : vec2D) {                                                     \\
    for (auto v : vec)                                                         \\
      cerr << setw(w) << v << " ";                                             \\
    cerr << endl;                                                              \\
  }
#define debugP(vec)                                                            \\
  for (auto v : vec)                                                           \\
    cerr << "(" << v.first << "," << v.second << ") ";                         \\
  cerr << endl;
#define debug2DP(vec2D)                                                        \\
  for (auto vec : vec2D) {                                                     \\
    for (auto v : vec)                                                         \\
      cerr << "(" << v.first << "," << v.second << ") ";                       \\
    cerr << endl;                                                              \\
  }

const char en = '\\n';

typedef long long ll;

constexpr ll INF = 1000000000; // 10^10
// const ll MOD = 998244353;
constexpr ll MOD = 1000000007;

typedef vector<vector<int>> Graph;
void add_edge(Graph &g, int from, int to, bool isDirected) {
  g[from].push_back(to);
  if (!isDirected)
    g[to].push_back(from);
}

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

int dx[4] = {0, 1, 0, -1}, dy[4] = {1, 0, -1, 0};

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

  ll n, k;
  cin >> n >> k;

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

  vector<vector<ll>> score(n, vector<ll>(n + 1, 0));
  vector<ll> loopCnt(n);

  rep(start, n) {
    ll sumi = 0;
    ll now = start, moveCnt = 0;
    while (p[now] != start) {
      sumi += c[p[now]];
      now = p[now];
      moveCnt++;
      score[start][moveCnt] = sumi;
    }
    ++moveCnt;
    loopCnt[start] = moveCnt;
    score[start][moveCnt] = sumi + c[p[now]];
  }

  // debug2D(score,3)

  ll ans = 0;
  rep(i, n) if (chmin(ans, c[i])) break;
  rep(start, n) {
    ll ansi = score[start][1];
    rep1(i, min(loopCnt[start], k)) {
      chmax(ansi, ((k - i) / loopCnt[start]) *
                          max(0LL, score[start][loopCnt[start]]) +
                      score[start][i]);
    }

    // cerr << "start = " << start << " : " << ansi << endl;

    chmax(ans, ansi);
  }
  cout << ans << endl;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cmath>
#include <iomanip> //double精度 setprecision
#include <iostream>
#include <map>
#include <numeric> //lcm
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>

// #include <atcoder/all>

using namespace std;
// using namespace atcoder;

#define rep(i, n) for (int i = 0; i < (n); ++i)
#define rrep(i, n) for (int i = n - 1; i >= 0; --i)
#define rep1(i, n) for (int i = 1; i <= (n); ++i)
#define rrep1(i, n) for (int i = (n); i >= 1; --i)
#define REP(i, n, m) for (int i = (n); i < (m); ++i)
#define all(vec) (vec).begin(), (vec).end()

#define debug(vec)                                                             \\
  for (auto v : vec)                                                           \\
    cerr << v << " ";                                                          \\
  cerr << endl;
#define debug2D(vec2D, w)                                                      \\
  for (auto vec : vec2D) {                                                     \\
    for (auto v : vec)                                                         \\
      cerr << setw(w) << v << " ";                                             \\
    cerr << endl;                                                              \\
  }
#define debugP(vec)                                                            \\
  for (auto v : vec)                                                           \\
    cerr << "(" << v.first << "," << v.second << ") ";                         \\
  cerr << endl;
#define debug2DP(vec2D)                                                        \\
  for (auto vec : vec2D) {                                                     \\
    for (auto v : vec)                                                         \\
      cerr << "(" << v.first << "," << v.second << ") ";                       \\
    cerr << endl;                                                              \\
  }

const char en = '\\n';

typedef long long ll;

constexpr ll INF = 1000000000; // 10^10
// const ll MOD = 998244353;
constexpr ll MOD = 1000000007;

typedef vector<vector<int>> Graph;
void add_edge(Graph &g, int from, int to, bool isDirected) {
  g[from].push_back(to);
  if (!isDirected)
    g[to].push_back(from);
}

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

int dx[4] = {0, 1, 0, -1}, dy[4] = {1, 0, -1, 0};

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

  ll n, k;
  cin >> n >> k;

  vector<ll> p(n), c(n);
  bool allneg = true;
  rep(i, n) {
    cin >> p[i];
    p[i]--;
  }
  rep(i, n) {
    cin >> c[i];
    allneg = allneg && (c[i] < 0);
  }

  if (allneg) {
    ll ans = c[0];
    rep(i, n) chmax(ans, c[0]);
    cout << ans << endl;
    return 0;
  }

  vector<vector<ll>> score(n, vector<ll>(n + 1, 0));
  vector<ll> loopCnt(n);

  rep(start, n) {
    ll sumi = 0;
    ll now = start, moveCnt = 0;
    while (p[now] != start) {
      sumi += c[p[now]];
      now = p[now];
      moveCnt++;
      score[start][moveCnt] = sumi;
    }
    ++moveCnt;
    loopCnt[start] = moveCnt;
    score[start][moveCnt] = sumi + c[p[now]];
  }

  debug2D(score, 3)

      ll ans = 0;
  rep(start, n) {
    ll ansi = 0, lastscore = 0;
    if (k >= loopCnt[start] && score[start][loopCnt[start]] > 0) {
      ansi += (k / loopCnt[start]) * score[start][loopCnt[start]];
      rep(i, k % loopCnt[start]) chmax(lastscore, score[start][i + 1]);
    } else {
      rep(i, min(n + 1, k + 1)) chmax(lastscore, score[start][i]);
    }

    ansi += lastscore;

    cerr << "start = " << start << " : " << ansi << endl;
    chmax(ans, ansi);
  }
  cout << ans << endl;
}