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

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

struct UnionFind {
  vector<int> par;

  UnionFind(int n) : par(n, -1) {}
  void init(int n) { par.assign(n, -1); }

  int root(int x) {
    if (par[x] < 0)
      return x;
    else
      return par[x] = root(par[x]);
  }

  bool issame(int x, int y) { return root(x) == root(y); }

  bool merge(int x, int y) {
    x = root(x);
    y = root(y);
    if (x == y)
      return false;
    if (par[x] > par[y])
      swap(x, y); // merge technique
    par[x] += par[y];
    par[y] = x;
    return true;
  }

  int size(int x) { return -par[root(x)]; }
};

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> P(N);
  vector<int> C(N);
  for (int i = 0; i < N; ++i) {
    cin >> P[i];
    P[i]--;
  }
  for (int i = 0; i < N; ++i) {
    cin >> C[i];
  }
  UnionFind uf(N);
  for (int n = 0; n < N; ++n) {
    uf.merge(n, P[n]);
  }
  map<int, long long> cost;
  for (int n = 0; n < N; ++n) {
    cost[uf.root(n)] += C[n];
  }

  long long ans = -1e18;
  for (int i = 0; i < N; ++i) {
    long long isize = uf.size(i);
    long long icost = cost[uf.root(i)];
    long long anstmp = 0;
    long long tmpK = K;
    if (tmpK >= 2 * isize) {
      long long n = tmpK / isize - 1;
      if (icost > 0) {
        anstmp += n * icost;
      }
      tmpK -= n * isize;
    }
    int now = i;
    long long maxsub = C[now];
    long long anssub = 0;
    for (int j = 0; j < tmpK; ++j) {
      // cout << now << "_" << C[now] << endl;
      anssub += C[now];
      now = P[now];
      maxsub = max(maxsub, anssub);
    }
    // cout << maxsub<<"_"<<anstmp<<endl;
    ans = max(maxsub + anstmp, ans);
  }
  cout << ans << endl;
  return 0;
}

Error: TLE

Fixed submission:

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

struct UnionFind {
  vector<int> par;

  UnionFind(int n) : par(n, -1) {}
  void init(int n) { par.assign(n, -1); }

  int root(int x) {
    if (par[x] < 0)
      return x;
    else
      return par[x] = root(par[x]);
  }

  bool issame(int x, int y) { return root(x) == root(y); }

  bool merge(int x, int y) {
    x = root(x);
    y = root(y);
    if (x == y)
      return false;
    if (par[x] > par[y])
      swap(x, y); // merge technique
    par[x] += par[y];
    par[y] = x;
    return true;
  }

  int size(int x) { return -par[root(x)]; }
};

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> P(N);
  vector<int> C(N);
  for (int i = 0; i < N; ++i) {
    cin >> P[i];
    P[i]--;
  }
  for (int i = 0; i < N; ++i) {
    cin >> C[i];
  }
  UnionFind uf(N);
  for (int n = 0; n < N; ++n) {
    uf.merge(n, P[n]);
  }
  map<int, long long> cost;
  for (int n = 0; n < N; ++n) {
    cost[uf.root(n)] += C[n];
  }

  long long ans = -1e18;
  for (int i = 0; i < N; ++i) {
    long long isize = uf.size(i);
    long long icost = cost[uf.root(i)];
    long long anstmp = 0;
    long long tmpK = K;
    if (tmpK >= 2 * isize) {
      if (icost > 0) {
        long long n = tmpK / isize - 1;
        anstmp += n * icost;
        tmpK -= n * isize;
      }
    }
    int now = i;
    long long maxsub = C[now];
    long long anssub = 0;
    for (int j = 0; j < tmpK; ++j) {
      // cout << now << "_" << C[now] << endl;
      anssub += C[now];
      now = P[now];
      maxsub = max(maxsub, anssub);
    }
    // cout << maxsub<<"_"<<anstmp<<endl;
    ans = max(maxsub + anstmp, ans);
  }
  cout << ans << endl;
  return 0;
}