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


// Problem :
// 			D - Moving Piece
//
//
// Contest : AtCoder - AtCoder Beginner Contest 175
// URL : https://atcoder.jp/contests/abc175/tasks/abc175_d
// Memory Limit : 1024 MB
// Time Limit : 3000 ms
// Powered by CP Editor (https://github.com/coder3101/cp-editor)

// $%U%$
// Mere bas ka nahi hai
// still i'm trying
// apna time aayega
//
//        .--------------.
//        | Try First One|
//        '--------------'
//                |     .--------------.
//                |     |              |
//                V     V              |
//              .--------------.       |
//              |      AC.     |<---.  |
//              '--------------'    |  |
//              (True)|  |(False)   |  |
//           .--------'  |          |  |
//           |           V          |  |
//           |  .--------------.    |  |
//           |  |   Try Again  |----'  |
//           |  '--------------'       |
//           |                         |
//           |  .--------------.       |
//           '->| Try Next One |-------'
//              '--------------'
//
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
ll power(ll x, ll y, ll m) {
  if (y == 0)
    return 1;
  ll p = power(x, y / 2, m) % m;
  p = (p * p) % m;
  return (y % 2 == 0) ? p : (x * p) % m;
}
ll nCr(ll n, ll r, ll m) {
  if (r > n)
    return 0;
  ll a = 1, b = 1, i;
  for (i = 0; i < r; i++) {
    a = (a * n) % m;
    --n;
  }
  while (r) {
    b = (b * r) % m;
    --r;
  }
  return (a * power(b, m - 2, m)) % m;
}
#define Test                                                                   \\
  int t;                                                                       \\
  cin >> t;                                                                    \\
  while (t--)
#define endl '\\n'
void solve() {
  ll n, k;
  cin >> n >> k;
  vector<ll> per(n);
  for (int i = 0; i < n; ++i) {
    cin >> per[i];
    per[i]--;
  }
  vector<ll> v(n);
  for (int i = 0; i < n; ++i)
    cin >> v[i];
  ll lgn = 40;
  vector<vector<ll>> bin(n, vector<ll>(lgn)), ans(n, vector<ll>(lgn));
  for (int i = 0; i < lgn; ++i) {
    for (int j = 0; j < n; ++j) {
      if (i == 0) {
        bin[j][i] = per[j];
        ans[j][i] = v[per[j]];
      } else {
        bin[j][i] = bin[bin[j][i - 1]][i - 1];
        ans[j][i] = ans[j][i - 1] + ans[bin[j][i - 1]][i - 1];
      }
    }
  }
  vector<ll> p;
  ll temp = 1;
  while (temp <= (ll)1e9) {
    p.push_back(temp);
    temp *= 2;
  }
  ll fans = -(ll)1e10;
  for (ll i = 0; i < n; ++i) {
    ll tempans = 0;
    ll a = i;
    for (ll take = 1; take <= k && take <= 2 * n; ++take) {
      a = per[a];
      tempans += v[a];
      fans = max(fans, tempans);
    }
    tempans = 0;
    if (k - 2 * n <= 0)
      continue;
    ll tk = k - 2 * n;
    a = i;
    for (int j = (int)p.size() - 1; j >= 0; --j) {
      if (tk >= p[j]) {
        tempans += ans[a][j];
        tk -= p[j];
        a = bin[a][j];
      }
    }
    fans = max(fans, tempans);
    for (ll take = 0; take < 2 * n; ++take) {
      a = per[a];
      tempans += v[a];
      fans = max(fans, tempans);
      // cout << tempans << " " << k+take+1 << " " << i << endl;
    }
    // cout << i << endl;
  }
  cout << fans << endl;
}
int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  // Test
  solve();
  return 0;
}


Error: TLE

Fixed submission:


// Problem :
// 			D - Moving Piece
//
//
// Contest : AtCoder - AtCoder Beginner Contest 175
// URL : https://atcoder.jp/contests/abc175/tasks/abc175_d
// Memory Limit : 1024 MB
// Time Limit : 3000 ms
// Powered by CP Editor (https://github.com/coder3101/cp-editor)

// $%U%$
// Mere bas ka nahi hai
// still i'm trying
// apna time aayega
//
//        .--------------.
//        | Try First One|
//        '--------------'
//                |     .--------------.
//                |     |              |
//                V     V              |
//              .--------------.       |
//              |      AC.     |<---.  |
//              '--------------'    |  |
//              (True)|  |(False)   |  |
//           .--------'  |          |  |
//           |           V          |  |
//           |  .--------------.    |  |
//           |  |   Try Again  |----'  |
//           |  '--------------'       |
//           |                         |
//           |  .--------------.       |
//           '->| Try Next One |-------'
//              '--------------'
//
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
ll power(ll x, ll y, ll m) {
  if (y == 0)
    return 1;
  ll p = power(x, y / 2, m) % m;
  p = (p * p) % m;
  return (y % 2 == 0) ? p : (x * p) % m;
}
ll nCr(ll n, ll r, ll m) {
  if (r > n)
    return 0;
  ll a = 1, b = 1, i;
  for (i = 0; i < r; i++) {
    a = (a * n) % m;
    --n;
  }
  while (r) {
    b = (b * r) % m;
    --r;
  }
  return (a * power(b, m - 2, m)) % m;
}
#define Test                                                                   \\
  int t;                                                                       \\
  cin >> t;                                                                    \\
  while (t--)
#define endl '\\n'
void solve() {
  ll n, k;
  cin >> n >> k;
  vector<ll> per(n);
  for (int i = 0; i < n; ++i) {
    cin >> per[i];
    per[i]--;
  }
  vector<ll> v(n);
  for (int i = 0; i < n; ++i)
    cin >> v[i];
  ll lgn = 40;
  vector<vector<ll>> bin(n, vector<ll>(lgn)), ans(n, vector<ll>(lgn));
  for (int i = 0; i < lgn; ++i) {
    for (int j = 0; j < n; ++j) {
      if (i == 0) {
        bin[j][i] = per[j];
        ans[j][i] = v[per[j]];
      } else {
        bin[j][i] = bin[bin[j][i - 1]][i - 1];
        ans[j][i] = ans[j][i - 1] + ans[bin[j][i - 1]][i - 1];
      }
    }
  }
  vector<ll> p;
  ll temp = 1;
  while (temp <= (ll)1e9) {
    p.push_back(temp);
    temp *= 2;
  }
  ll fans = -(ll)1e10;
  for (ll i = 0; i < n; ++i) {
    for (ll take = 1; take <= k && take <= 2 * n; ++take) {
      ll tans = 0;
      ll tk = take;
      ll a = i;
      for (int j = (int)p.size() - 1; j >= 0; --j) {
        if (tk >= p[j]) {
          tans += ans[a][j];
          tk -= p[j];
          a = bin[a][j];
        }
      }
      // cout << tans << " ";
      fans = max(fans, tans);
      // if(tans==(ll)29507023469)
      // cout << tans << " " << i << " " << take << endl;
    }
    for (ll take = k; take >= k - 2 * n - 1 && take >= 1; --take) {
      ll tans = 0;
      ll tk = take;
      ll a = i;
      for (int j = (int)p.size() - 1; j >= 0; --j) {
        if (tk >= p[j]) {
          tans += ans[a][j];
          tk -= p[j];
          a = bin[a][j];
        }
      }
      fans = max(fans, tans);
    }
  }
  cout << fans << endl;
}
int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  // Test
  solve();
  return 0;
}