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

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

#define rep(i, n) for (int i = 0; i < n; i++)
#define Rep(i, n) for (int i = 1; i <= n; i++)
#define sz(x) int(x.size())
#define all(v) v.begin(), v.end()
#define rall(v) v.rbegin(), v.rend()
#define YesorNo(a) printf(a ? "Yes\\n" : "No\\n")
#define endl '\\n'
#define fi first
#define se second
using ll = long long;
using P = pair<int, int>;
using Pl = pair<ll, ll>;
template <class T> using V = vector<T>;
const int dx[] = {0, 1, 0, -1, 1, 1, -1, -1};
const int dy[] = {1, 0, -1, 0, 1, -1, -1, 1};
const int inf = (1 << 30) - 1;
const ll infll = (1LL << 62) - 1;
ll ceil(const ll &a, const ll &b) { return ((a) + (b)-1) / b; }
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

struct UnionFind { // UnionFindTree
  vector<int> d;
  UnionFind(int n) : d(n, -1) {}
  int find(int x) {
    if (d[x] < 0)
      return x;
    return d[x] = find(d[x]);
  }
  bool unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return false;
    if (-d[x] < -d[y])
      swap(x, y);
    d[x] += d[y];
    d[y] = x;
    return true;
  }
  bool same(int x, int y) { return find(x) == find(y); }
  int size(int x) { return -d[find(x)]; }
};

int main() {
  int n, k;
  cin >> n >> k;
  V<int> p(n), c(n);
  rep(i, n) cin >> p[i];
  rep(i, n) cin >> c[i];
  map<int, ll> mp;
  UnionFind uf(n);
  rep(i, n) { uf.unite(i, p[i] - 1); }
  rep(i, n) { mp[uf.find(i)] += c[i]; }
  ll ans = -infll;
  rep(i, n) { // スタート位置
    int roop = uf.size(i);
    int times = (mp[uf.find(i)] > 0 ? (k - 1) / roop : 0);
    int over = max(min(roop, k - roop * times), 1);
    ll tmp = 0;
    ll s = -infll;
    int idx = i;
    rep(j, over) {
      tmp += c[idx];
      chmax(s, tmp);
      idx = p[idx] - 1;
    }
    s += times * mp[uf.find(i)];
    chmax(ans, s);
  }
  cout << ans << endl;
}

Error: TLE

Fixed submission:

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

#define rep(i, n) for (int i = 0; i < n; i++)
#define Rep(i, n) for (int i = 1; i <= n; i++)
#define sz(x) int(x.size())
#define all(v) v.begin(), v.end()
#define rall(v) v.rbegin(), v.rend()
#define YesorNo(a) printf(a ? "Yes\\n" : "No\\n")
#define endl '\\n'
#define fi first
#define se second
using ll = long long;
using P = pair<int, int>;
using Pl = pair<ll, ll>;
template <class T> using V = vector<T>;
const int dx[] = {0, 1, 0, -1, 1, 1, -1, -1};
const int dy[] = {1, 0, -1, 0, 1, -1, -1, 1};
const int inf = (1 << 30) - 1;
const ll infll = (1LL << 62) - 1;
ll ceil(const ll &a, const ll &b) { return ((a) + (b)-1) / b; }
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

struct UnionFind { // UnionFindTree
  vector<int> d;
  UnionFind(int n) : d(n, -1) {}
  int find(int x) {
    if (d[x] < 0)
      return x;
    return d[x] = find(d[x]);
  }
  bool unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return false;
    if (-d[x] < -d[y])
      swap(x, y);
    d[x] += d[y];
    d[y] = x;
    return true;
  }
  bool same(int x, int y) { return find(x) == find(y); }
  int size(int x) { return -d[find(x)]; }
};

int main() {
  int n, k;
  cin >> n >> k;
  V<int> p(n), c(n);
  rep(i, n) cin >> p[i];
  rep(i, n) cin >> c[i];
  map<int, ll> mp;
  UnionFind uf(n);
  rep(i, n) { uf.unite(i, p[i] - 1); }
  rep(i, n) { mp[uf.find(i)] += c[i]; }
  ll ans = -infll;
  rep(i, n) { // スタート位置
    int roop = uf.size(i);
    int times = (mp[uf.find(i)] > 0 ? (k - 1) / roop : 0);
    int over = max(min(k, k - roop * times), 1);
    ll tmp = 0;
    ll s = -infll;
    int idx = i;
    rep(j, over) {
      tmp += c[idx];
      chmax(s, tmp);
      idx = p[idx] - 1;
    }
    s += times * mp[uf.find(i)];
    chmax(ans, s);
  }
  cout << ans << endl;
}