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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <vector>

#ifdef DEBUG
#include "library/Utility/debug.cpp"
#else
#define debug(...)
#endif

#define rep(i, n) for (int i = 0; i < (n); ++i)
#define EL '\\n'
#define print(i) std::cout << (i) << '\\n'
#define all(v) (v).begin(), (v).end()
using lnt = long long;
struct FIO {
  FIO() {
    std::cin.tie(0);
    std::cout.tie(0);
    std::ios_base::sync_with_stdio(0);
    std::cout << std::fixed << std::setprecision(15);
  }
} fIO;
/*-*/

struct UnionFind {
  int n;
  std::vector<int> par, size;
  UnionFind(int sz) : n(sz), par(sz, 0), size(sz, 1) {
    std::iota(par.begin(), par.end(), 0);
  }
  int find(int x) {
    if (par[x] == x)
      return x;
    return par[x] = find(par[x]);
  }
  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;
    if (size[x] < size[y])
      std::swap(x, y);
    size[x] += size[y];
    par[y] = x;
  }
};

int main() {
  lnt m, k;
  std::cin >> m >> k;
  std::vector<lnt> p(m), c(m);
  rep(i, m) std::cin >> p[i];
  rep(i, m) std::cin >> c[i];
  lnt max = -3e18;
  std::map<lnt, lnt> map;
  UnionFind uf(m);
  rep(i, m) { uf.unite(i, p[i] - 1); }
  rep(ii, m) {
    if (map[uf.find(ii)])
      continue;
    map[uf.find(ii)] = 1;
    lnt n = uf.size[uf.find(ii)];
    std::vector<lnt> dp(n + 1, 0);
    {
      int i = ii;
      rep(j, n) {
        dp[j + 1] = dp[j] + c[i];
        i = p[i] - 1;
      }
    }
    rep(i, n) {
      for (int j = i; j < n; j++) {
        if (j - i + 1 <= k)
          max = std::max(max, dp[j + 1] - dp[i]);
        if (n - i + j + 1 <= k) {
          max = std::max(max, dp[n] - dp[i] + dp[j + 1]);
          max = std::max(max, dp[n] - dp[i] + dp[j + 1] +
                                  dp[n] * ((k - (n - i + j + 1)) / n));
        }
        if (n - j + i + 1 <= k) {
          max = std::max(max, dp[n] - dp[j] + dp[i + 1]);
          max = std::max(max, dp[n] - dp[j] + dp[i + 1] +
                                  dp[n] * ((k - (n - j + i + 1)) / n));
        }
      }
    }
  }
  print(max);
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <vector>

#ifdef DEBUG
#include "library/Utility/debug.cpp"
#else
#define debug(...)
#endif

#define rep(i, n) for (int i = 0; i < (n); ++i)
#define EL '\\n'
#define print(i) std::cout << (i) << '\\n'
#define all(v) (v).begin(), (v).end()
using lnt = long long;
struct FIO {
  FIO() {
    std::cin.tie(0);
    std::cout.tie(0);
    std::ios_base::sync_with_stdio(0);
    std::cout << std::fixed << std::setprecision(15);
  }
} fIO;
/*-*/

struct UnionFind {
  int n;
  std::vector<int> par, size;
  UnionFind(int sz) : n(sz), par(sz, 0), size(sz, 1) {
    std::iota(par.begin(), par.end(), 0);
  }
  int find(int x) {
    if (par[x] == x)
      return x;
    return par[x] = find(par[x]);
  }
  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;
    if (size[x] < size[y])
      std::swap(x, y);
    size[x] += size[y];
    par[y] = x;
  }
};

int main() {
  lnt m, k;
  std::cin >> m >> k;
  std::vector<lnt> p(m), c(m);
  rep(i, m) std::cin >> p[i];
  rep(i, m) std::cin >> c[i];
  lnt max = -3e18;
  std::map<lnt, lnt> map;
  UnionFind uf(m);
  rep(i, m) { uf.unite(i, p[i] - 1); }
  rep(ii, m) {
    if (map[uf.find(ii)])
      continue;
    lnt n = uf.size[uf.find(ii)];
    std::vector<lnt> dp(n + 1, 0);
    {
      int i = ii;
      rep(j, n) {
        dp[j + 1] = dp[j] + c[i];
        i = p[i] - 1;
      }
    }
    rep(i, n) {
      for (int j = i; j < n; j++) {
        if (j - i + 1 <= k)
          max = std::max(max, dp[j + 1] - dp[i]);
        if (n - i + j + 1 <= k) {
          max = std::max(max, dp[n] - dp[i] + dp[j + 1]);
          max = std::max(max, dp[n] - dp[i] + dp[j + 1] +
                                  dp[n] * ((k - (n - i + j + 1)) / n));
        }
        if (n - j + i + 1 <= k) {
          max = std::max(max, dp[n] - dp[j] + dp[i + 1]);
          max = std::max(max, dp[n] - dp[j] + dp[i + 1] +
                                  dp[n] * ((k - (n - j + i + 1)) / n));
        }
      }
    }
  }
  print(max);
}