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

// Konrad Paluszek,University of Warsaw(former XIV LO Staszic)
// #STAY AT HOME
#ifndef LOCAL
#pragma GCC optimize("O3")
#endif
#define TIME (chrono::steady_clock::now().time_since_epoch().count())
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#define xfm(a, b) a##b
#define xwb(a, b) xfm(a, b)
#define _ xwb(nvj, __LINE__)
#define __ xwb(kjy, __LINE__)
#define ___ xwb(cjm, __LINE__)
#define REP(i, n) for (urs(n) i = 0; i < (n); ++i)
#define UNTIL(t) while (clock() < (t)*CLOCKS_PER_SEC)
#define PER(r...) for (bool _ = 1; _ || next_permutation(r); _ = false)
#define ALL(r) (r).begin(), (r).end()
#define RALL(r) (r).rbegin(), (r).rend()
#define FS(r) r.first, r.second
#define SF(r) r.second, r.first
#define M0(r) memset(r, 0, sizeof(r))
#define sim template <class c
#define ros return o
#define rans return ans
#define forbits(i, m)                                                          \\
  if (m)                                                                       \\
    for (urs(m) i = ctz(m), i##nvj = m; i##nvj;                                \\
         i##nvj ^= ((urs(m))1 << i), i = ctz(i##nvj))
#define fordbits(i, m)                                                         \\
  if (m)                                                                       \\
    for (urs(m) i = 8 * sizeof(m) - clz(m) - 1, i##nxd = m; i##nxd;            \\
         i##nxd ^= ((urs(m))1 << i), i = 8 * sizeof(m) - clz(i##nxd) - 1)
#define ksets(t, m, k, n)                                                      \\
  for (t m = (((t)1 << (k)) - 1); m < ((t)1 << (n)); m = nux(m))
#define urs(r...) typename decay<decltype(r)>::type
#define hur(f, g, r)                                                           \\
  sim > int f(c a) {                                                           \\
    if (sizeof(c) == 16)                                                       \\
      return r;                                                                \\
    if (sizeof(c) == 8)                                                        \\
      return g##ll(a);                                                         \\
    return g(a);                                                               \\
  }
#define pwq(t, i)                                                              \\
  int clz(t x) { return clz<int>(x) - i; }
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wconversion"
using namespace __gnu_pbds;
using namespace std;
using ll = long long;
using ld = long double;
using ull = unsigned long long;
using vi = vector<int>;
using vll = vector<ll>;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using vpii = vector<pii>;
using spii = set<pii>;
using mii = map<int, int>;
using unt = unsigned int;
sim > using min_queue = priority_queue<c, vector<c>, greater<c>>;
sim, class b,
    class cmp =
        less<c> > using ordered_map =
            tree<c, b, cmp, rb_tree_tag, tree_order_statistics_node_update>;
sim, class cmp = less < c >> using ordered_set = ordered_map<c, null_type, cmp>;
hur(popc, __builtin_popcount, popc<ull>(a) + popc<ull>(a >> 64))
    hur(ctz, __builtin_ctz, (ull)a ? ctz<ull>(a) : 64 + ctz<ull>(a >> 64))
        hur(clz, __builtin_clz, a >> 64 ? clz<ull>(a >> 64) : 64 + clz<ull>(a))
            pwq(short, 16) pwq(uint16_t, 16) pwq(char, 24) pwq(int8_t, 24)
                pwq(uint8_t, 24) sim,
    class N > bool mini(c &o, const N &h) {
  if (o > h)
    ros = h, 1;
  return 0;
}
sim, class N > bool maxi(c &o, const N &h) {
  if (o < h)
    ros = h, 1;
  return 0;
}
sim, class n > using gyv = c;
#if defined(LOCAL) // || defined(LOCAL2)
#include </home/kjp/headers/debuglib.hpp>
#else
#define loc(...)
#define onl(r...) r
#define debug(...)
#define print_stack(...)
#define mark_stack(...)
#define set_pre(...)
#define reg_it(...)
#define def_op(...)                                                            \\
  struct _ {};
#define mask_set(...)
#if !defined(LOCAL) && !defined(LOCAL2)
#define exit my_exit
void my_exit(int x) {
  fflush(stdout);
  _Exit(x);
}
#endif
#endif
#define next nexT
#define prev preV
#define tree trEE
#define left lefT
#define right righT
#define div diV
#define y1 y_1
#define pow                                                                    \\
  do                                                                           \\
  not use cmath pow unless you know what you are doing
ull mix(ull o) {
  o += 0x9e3779b97f4a7c15;
  o = (o ^ (o >> 30)) * 0xbf58476d1ce4e5b9;
  o = (o ^ (o >> 27)) * 0x94d049bb133111eb;
  ros ^ (o >> 31);
}
ull SALT = 0x7a14a4b0881ebf9, tqu = 0x7a14a4b0881ebf9;
ull my_rand() { return tqu = mix(tqu); }
void my_srand(ull x) { SALT = tqu = x; }
const int inf = 1023400000;
const ll llinf = 1234567890000000000ll;
ll fix(ll o, ll m) {
  o %= m;
  if (o < 0)
    o += m;
  ros;
}
#define rand my_rand
#define srand my_srand
#define random_shuffle(r...)                                                   \\
  random_shuffle(r, [](int _) { return my_rand() % _; })
sim > inline c nux(c m) {
  if (!m)
    return numeric_limits<c>::max();
  c A = m & -m;
  c B = ~((A - 1) ^ m);
  c C = B & -B;
  c D = (C >> (1 + ctz(A))) - 1;
  return C | (m & ~(C - 1)) | D;
}
__attribute__((no_sanitize_undefined)) ll mul(ll a, ll b, ll m) {
  ll q = (ll)(a * (ld)b / m);
  ll o = a * b - q * m;
  o %= m;
  if (o < 0)
    o += m;
  ros;
}
sim > void unq(c &x) { x.resize(unique(ALL(x)) - x.begin()); }
#pragma GCC diagnostic pop
#if ((ULONG_MAX) != (UINT_MAX))
namespace std {
template <> struct is_signed<__int128> : public true_type {};
} // namespace std
#endif
sim, class d > typename common_type<c, d>::type floor_div(c a, d b) {
  static_assert(is_signed<c>::value == is_signed<d>::value,
                "using floor_div with different signedness");
  if (b < 0)
    b = -b, a = -a;
  return a / b - (a % b < 0);
}
sim, class d > typename common_type<c, d>::type ceil_div(c a, d b) {
  static_assert(is_signed<c>::value == is_signed<d>::value,
                "using ceil_div with different signedness");
  if (b < 0)
    b = -b, a = -a;
  return a / b + (a % b > 0);
}
sim > struct REV {
  using value_type = typename c::value_type;
  c &x;
  using it = typename c::reverse_iterator;
  it begin() { return x.rbegin(); }
  it end() { return x.rend(); }
};
sim > struct CREV {
  using value_type = typename c::value_type;
  const c &x;
  using it = typename c::const_reverse_iterator;
  it begin() { return x.rbegin(); }
  it end() { return x.rend(); }
};
sim > REV<c> reversed(c &x) { return REV<c>{x}; }
sim > CREV<c> reversed(const c &x) { return CREV<c>{x}; }
#define done(r...) exit(0 * printf(r))
#if defined(LOCAL) || defined(LOCAL2)
void __tmi() {
  cerr << setprecision(6) << fixed
       << "total time: " << clock() / (ld)CLOCKS_PER_SEC << "s" << endl;
}
int _ = (atexit(__tmi), 0);
#endif
// #STAY AT HOME
ll get_one(const vi &vals, int lim, bool can_empty) {
  int n = vals.size();
  assert(lim < n);
  vll sums(2 * n + 1);
  REP(i, 2 * n) sums[i + 1] = sums[i] + vals[i % n];
  multiset<ll> seen;
  ll ans = can_empty ? 0 : numeric_limits<ll>::min();
  for (int i = 0; i <= 2 * n; ++i) {
    if (!seen.empty())
      maxi(ans, sums[i] - *seen.begin());
    seen.insert(sums[i]);
    int l = i - lim;
    if (l >= 0)
      seen.erase(seen.find(sums[l]));
  }
  debug(imie(vals), imie(lim), imie(ans), imie(can_empty));
  rans;
}
void solve() {
  int n, k;
  scanf("%d%d", &n, &k);
  vi p(n + 1), c(n + 1);
  for (int i = 1; i <= n; ++i)
    scanf("%d", &p[i]);
  for (int i = 1; i <= n; ++i)
    scanf("%d", &c[i]);
  vi did(n + 1);
  ll ans = numeric_limits<ll>::min();
  for (int i = 1; i <= n; ++i) {
    if (!did[i]) {
      int x = i;
      vi ord;
      do {
        ord.push_back(c[x]);
        did[x] = true;
        x = p[x];
      } while (x != i);
      debug(imie(ord));
      ll su = 0;
      int le = ord.size();
      for (int x : ord)
        su += x;
      maxi(ans, get_one(ord, min(k, le - 1), false));
      int full = k / le, rest = k % le;
      maxi(ans, get_one(ord, rest, full >= 1) + su * 1ll * full);
      if (full >= 1)
        maxi(ans, get_one(ord, le - 1, full > 1) + su * 1ll * (full - 1));
    }
  }
  printf("%lld\\n", ans);
}
int main() {
  // unt seed = TIME; debug(imie(seed));srand(seed);
  int t = 1;
  // scanf("%d", &t);
  REP(_, t) solve();
  exit(0);
}
// #STAY AT HOME


Error: TLE

Fixed submission:

/*created by tanishk gupta*/
#include <bits/stdc++.h>
using namespace std;
#define rep(i, a, n) for (int i = a; i < n; i++)
#define per(i, a, n) for (int i = n - 1; i >= a; i--)
#define repl(i, a, n) for (long long int i = a; i < n; i++)
#define pb push_back
#define mp make_pair
#define all(x) x.begin(), x.end()
#define fi first
#define se second
#define SZ(x) ((int)(x).size())
#define IOS                                                                    \\
  ios::sync_with_stdio(0);                                                     \\
  cin.tie(0);                                                                  \\
  cout.tie(0);
typedef vector<int> VI;
typedef long long int ll;
typedef unsigned long long ull;
typedef pair<int, int> PII;
typedef double db;
mt19937 mrand(random_device{}());
const ll MOD = 1e9 + 7;
const ll mod = 1e9 + 7;
// const llint MOD;
const ll N = (ll)(1e6 + 1);
int rnd(int x) { return mrand() % x; }
ll powmod(ll a, ll b) {
  ll res = 1;
  a %= mod;
  assert(b >= 0);
  for (; b; b >>= 1) {
    if (b & 1)
      res = res * a % mod;
    a = a * a % mod;
  }
  return res;
}
ll gcd(ll a, ll b) { return b ? gcd(b, a % b) : a; }
ll add(ll a, ll b) { return ((a % MOD) + (b % MOD)) % MOD; }
ll mul(ll a, ll b) { return ((a % MOD) * (b % MOD)) % MOD; }
ll sub(ll a, ll b) { return ((a % MOD) - (b % MOD) + MOD) % MOD; }
ll binpow(ll x, ll y) {
  ll z = 1;
  while (y) {
    if (y & 1)
      z = mul(z, x);
    x = mul(x, x);
    y >>= 1;
  }
  return z;
}
ll inv(int x) { return binpow(x, MOD - 2); }
ll divide(ll x, ll y) { return mul(x, inv(y)); }
ll fact[N];
void precalc() {
  fact[0] = 1;
  for (ll i = 1; i < N; i++)
    fact[i] = mul(fact[i - 1], i);
}
ll C(ll n, ll k) { return divide(fact[n], mul(fact[k], fact[n - k])); }
// bool prime[N];
// void sieve() {for(int i=2;i<N;i++) {prime[i]=true;}for (int p=2; p*p<=N;
// p++){if (prime[p] == true){for (int i=p*p; i<=N; i += p)prime[i] = false;}}}
ll legendpow(ll n, ll p) {
  ll total = 0;
  while (n) {
    n /= p;
    total += n;
  }
  return total;
}

// head==================================================================================
// ll dp[1024][1024];
int main() {

  int n, k;
  cin >> n >> k;
  int p[n + 1];
  int c[n + 1];
  // int c[]
  // ll ans=0;
  rep(i, 1, n + 1) cin >> p[i];
  rep(i, 1, n + 1) cin >> c[i];
  ll tot = 0;
  rep(i, 1, n + 1) {
    int temp = 2;
    ll ans = c[p[i]];
    ll curr = c[p[i]];
    int index = p[i];
    while (temp <= k) { // temp++;
                        //  i=p[i];
      curr += c[p[index]];
      ans = max(ans, curr);
      temp++;
      index = p[index];
    }
    tot = max(ans, tot);
  }
  if (tot == 0) {
    tot = *max_element(c + 1, c + n);
    cout << tot;
  } else {

    cout << tot << endl;
  }
}