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

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

using ll = long long;
constexpr ll mod = 1e9 + 7;
#define endl '\\n'

template <typename A, typename B> string to_string(pair<A, B> p);

template <typename A, typename B, typename C>
string to_string(tuple<A, B, C> p);

template <typename A, typename B, typename C, typename D>
string to_string(tuple<A, B, C, D> p);

string to_string(const string &s) { return '"' + s + '"'; }

string to_string(const char c) { return string(1, c); }

string to_string(const char *s) { return to_string((string)s); }

string to_string(bool b) { return (b ? "true" : "false"); }

string to_string(vector<bool> v) {
  bool first = true;
  string res = "{";
  for (int i = 0; i < static_cast<int>(v.size()); i++) {
    if (!first) {
      res += ", ";
    }
    first = false;
    res += to_string(v[i]);
  }
  res += "}";
  return res;
}

template <size_t N> string to_string(bitset<N> v) {
  string res = "";
  for (int i = int(N) - 1; i >= 0; --i) {
    res += static_cast<char>('0' + v[i]);
  }
  return res;
}

template <typename A> string to_string(A v) {
  bool first = true;
  string res = "{";
  for (const auto &x : v) {
    if (!first) {
      res += ", ";
    }
    first = false;
    res += to_string(x);
  }
  res += "}";
  return res;
}

template <typename A, typename B> string to_string(pair<A, B> p) {
  return "(" + to_string(p.first) + ", " + to_string(p.second) + ")";
}

template <typename A, typename B, typename C>
string to_string(tuple<A, B, C> p) {
  return "(" + to_string(get<0>(p)) + ", " + to_string(get<1>(p)) + ", " +
         to_string(get<2>(p)) + ")";
}

template <typename A, typename B, typename C, typename D>
string to_string(tuple<A, B, C, D> p) {
  return "(" + to_string(get<0>(p)) + ", " + to_string(get<1>(p)) + ", " +
         to_string(get<2>(p)) + ", " + to_string(get<3>(p)) + ")";
}

void debug_out() { cerr << endl; }

template <typename Head, typename... Tail> void debug_out(Head H, Tail... T) {
  cerr << " " << to_string(H);
  debug_out(T...);
}

#ifdef COMPETITIVE_PROGRAMMING_LOCAL
#define dbg(...) cerr << "[" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#else
#define dbg(...)                                                               \\
  if (0) {                                                                     \\
  }
#endif

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

  ll N, K;
  cin >> N >> K;
  vector<ll> P(N), C(N);
  for (auto &_a : P)
    cin >> _a, --_a;
  for (auto &_a : C)
    cin >> _a;

  vector<ll> visited(N);
  vector<ll> sum(N);
  ll ans = *max_element(C.begin(), C.end());
  for (ll start = 0; start < N; ++start) {
    fill(visited.begin(), visited.end(), -1);
    fill(sum.begin(), sum.end(), 0);

    ll pos = start;
    ll score = 0;
    ll k = 0;

    while (k < K) {
      pos = P[pos];
      if (visited[pos] != -1) {
        // on the last cycle we don't need to perform full cycle
        if (score - sum[pos] > 0) {
          ll n_step_in_cycle = k - visited[pos];
          ll n_step_remain = K - k;
          ll n_cycle = n_step_remain / n_step_in_cycle;

          //          dbg(pos, visited[pos], sum[pos], score, n_step_remain, K,
          //          k);

          n_cycle = max(0LL, n_cycle - 1);
          k += n_cycle * n_step_in_cycle;
          score += n_cycle * (score - sum[pos]);
          ans = max(ans, score);
          while (k < K) {
            score += C[pos];
            ans = max(ans, score);
            ++k;
            pos = P[pos];
          }

        } else {
          // continue while positive, do one last loop in case start of cycle is
          // positive
          ll idx = 0;
          while (k < K && idx < 2 * N) {
            score += C[pos];
            ans = max(ans, score);
            pos = P[pos];
            ++idx;
            ++k;
          }
        }
        break;
      } else {
        sum[pos] = score;
        score += C[pos];
        ans = max(ans, score);
        visited[pos] = k;
      }
      ++k;
    }
  }
  cout << ans << endl;
}


Error: TLE

Fixed submission:

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

using ll = long long;
constexpr ll mod = 1e9 + 7;
#define endl '\\n'

template <typename A, typename B> string to_string(pair<A, B> p);

template <typename A, typename B, typename C>
string to_string(tuple<A, B, C> p);

template <typename A, typename B, typename C, typename D>
string to_string(tuple<A, B, C, D> p);

string to_string(const string &s) { return '"' + s + '"'; }

string to_string(const char c) { return string(1, c); }

string to_string(const char *s) { return to_string((string)s); }

string to_string(bool b) { return (b ? "true" : "false"); }

string to_string(vector<bool> v) {
  bool first = true;
  string res = "{";
  for (int i = 0; i < static_cast<int>(v.size()); i++) {
    if (!first) {
      res += ", ";
    }
    first = false;
    res += to_string(v[i]);
  }
  res += "}";
  return res;
}

template <size_t N> string to_string(bitset<N> v) {
  string res = "";
  for (int i = int(N) - 1; i >= 0; --i) {
    res += static_cast<char>('0' + v[i]);
  }
  return res;
}

template <typename A> string to_string(A v) {
  bool first = true;
  string res = "{";
  for (const auto &x : v) {
    if (!first) {
      res += ", ";
    }
    first = false;
    res += to_string(x);
  }
  res += "}";
  return res;
}

template <typename A, typename B> string to_string(pair<A, B> p) {
  return "(" + to_string(p.first) + ", " + to_string(p.second) + ")";
}

template <typename A, typename B, typename C>
string to_string(tuple<A, B, C> p) {
  return "(" + to_string(get<0>(p)) + ", " + to_string(get<1>(p)) + ", " +
         to_string(get<2>(p)) + ")";
}

template <typename A, typename B, typename C, typename D>
string to_string(tuple<A, B, C, D> p) {
  return "(" + to_string(get<0>(p)) + ", " + to_string(get<1>(p)) + ", " +
         to_string(get<2>(p)) + ", " + to_string(get<3>(p)) + ")";
}

void debug_out() { cerr << endl; }

template <typename Head, typename... Tail> void debug_out(Head H, Tail... T) {
  cerr << " " << to_string(H);
  debug_out(T...);
}

#ifdef COMPETITIVE_PROGRAMMING_LOCAL
#define dbg(...) cerr << "[" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#else
#define dbg(...)                                                               \\
  if (0) {                                                                     \\
  }
#endif

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

  ll N, K;
  cin >> N >> K;
  vector<ll> P(N), C(N);
  for (auto &_a : P)
    cin >> _a, --_a;
  for (auto &_a : C)
    cin >> _a;

  vector<ll> visited(N);
  vector<ll> sum(N);
  ll ans = *max_element(C.begin(), C.end());
  for (ll start = 0; start < N; ++start) {
    fill(visited.begin(), visited.end(), -1);
    fill(sum.begin(), sum.end(), 0);

    ll pos = start;
    ll score = 0;
    ll k = 0;

    while (k < K) {
      pos = P[pos];
      if (visited[pos] != -1) {
        // on the last cycle we don't need to perform full cycle
        if (score - sum[pos] > 0) {
          ll n_step_in_cycle = k - visited[pos];
          ll n_step_remain = K - k;
          ll n_cycle = n_step_remain / n_step_in_cycle;

          //          dbg(pos, visited[pos], sum[pos], score, n_step_remain, K,
          //          k);

          n_cycle = max(0LL, n_cycle - 1);
          k += n_cycle * n_step_in_cycle;
          score += n_cycle * (score - sum[pos]);
          ans = max(ans, score);
          while (k < K) {
            score += C[pos];
            ans = max(ans, score);
            ++k;
            pos = P[pos];
          }

        } else {
          // continue while positive, do one last loop in case start of cycle is
          // positive
          ll idx = 0;
          while (k < K && idx < 2 * N) {
            score += C[pos];
            ans = max(ans, score);
            pos = P[pos];
            ++idx;
            ++k;
          }
        }
      } else {
        sum[pos] = score;
        score += C[pos];
        ans = max(ans, score);
        visited[pos] = k;
      }
      ++k;
    }
  }
  cout << ans << endl;
}