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

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

#pragma region Macros
using ll = long long;
#define int ll
using pii = pair<int, int>;
using tiii = tuple<int, int, int>;
template <class T = ll> using V = vector<T>;
template <class T = ll> using VV = V<V<T>>;
#define IOS                                                                    \\
  ios::sync_with_stdio(false);                                                 \\
  cin.tie(0);                                                                  \\
  cout.tie(0);
#define FOR(i, l, r) for (ll i = (l); i < (r); ++i)
#define REP(i, n) FOR(i, 0, n)
#define REPS(i, n) FOR(i, 1, n + 1)
#define RFOR(i, l, r) for (ll i = (l); i >= (r); --i)
#define RREP(i, n) RFOR(i, n - 1, 0)
#define RREPS(i, n) RFOR(i, n, 1)
#define mp make_pair
#define mt make_tuple
#define pb push_back
#define eb emplace_back
#define all(x) (x).begin(), (x).end()
#define SORT(name) sort(name.begin(), name.end())
#define RSORT(name)                                                            \\
  SORT(name);                                                                  \\
  reverse(all(name));
#define ZERO(p) memset(p, 0, sizeof(p))
#define MINUS(p) memset(p, -1, sizeof(p))
inline void Yes(bool b = true) { cout << (b ? "Yes" : "No") << '\\n'; }
inline void YES(bool b = true) { cout << (b ? "YES" : "NO") << '\\n'; }
template <class T> inline void print(T x) { cout << x << '\\n'; }
template <typename T1, typename T2> inline void chmin(T1 &a, T2 b) {
  if (a > b)
    a = b;
}
template <typename T1, typename T2> inline void chmax(T1 &a, T2 b) {
  if (a < b)
    a = b;
}

const ll LLINF = (1LL << 60);
const int INF = (1LL << 30);
const double DINF = std::numeric_limits<double>::infinity();
#pragma endregion

#if 1
#define DBG(fmt, ...) printf(fmt, ##__VA_ARGS__)
#else
#define DBG(fmt, ...)
#endif

class UnionFind {
public:
  UnionFind() {}
  UnionFind(int size) { Init(size); }
  // size: Union-Find で扱いたい要素の数
  void Init(int size) { m_Data.assign(size, -1); }

  // x と y の属する集合を併合
  // 既に一緒な集合内にいるなら false を
  // 新たに併合されたなら true を返す
  bool Unite(int x, int y) {
    x = Find(x), y = Find(y);
    if (x == y) {
      return (false);
    }
    if (m_Data[x] > m_Data[y]) {
      swap(x, y);
    }
    m_Data[x] += m_Data[y];
    m_Data[y] = x;
    return (true);
  }

  // 要素 k が属する集合の根を返す
  int Find(int k) {
    if (m_Data[k] < 0) {
      return (k);
    }
    return (m_Data[k] = Find(m_Data[k]));
  }

  // 要素 k と同じ集合の要素数を返す
  int Size(int k) { return (-m_Data[Find(k)]); }

  // 要素のうち独立しているグループの数 O(NlogN)
  int GetGroupNum() {
    map<int, bool> groups;
    for (int i = 0; i < m_Data.size(); ++i) {
      groups[Find(i)] = true;
    }
    return static_cast<int>(groups.size());
  }

private:
  vector<int> m_Data;
};

#define SUBMIT

const int MOD = 1000000007;
const int MAX_N = 100010;

int N, K;
V<> P;
V<> C;

UnionFind uf;
V<bool> isUsed;
int ans = -LLINF;
V<> group; // グループ i の1周の点数

int dfs(int i) {
  if (isUsed[i]) {
    return 0;
  }
  isUsed[i] = true;

  uf.Unite(i, P[i]);
  return dfs(P[i]) + C[P[i]];
}

void dfs2(int i, int cur, int rem) {
  // if(isUsed[i]) { return; }
  if (rem <= 0) {
    return;
  }
  // isUsed[i] = true;

  cur += C[P[i]];
  chmax(ans, cur);
  dfs2(P[i], cur, rem - 1);
}

int solve() {
  REP(i, N) { chmax(ans, C[i]); }

  uf.Init(N);

  // group の作成
  group.assign(N, 0);
  isUsed.assign(N, false);
  REP(i, N) {
    if (isUsed[i]) {
      continue;
    }
    int val = dfs(i);
    int g_id = uf.Find(i);
    group[g_id] = val;
  }

  REP(i, N) {
    int g_id = uf.Find(i);
    // 1 周にかかるコスト
    int cost = uf.Size(i);
    // isUsed.assign(N, false);
    if (group[g_id] > 0) {
      // ぐるぐる回る
      int cand = (K / cost - 1) * group[g_id];
      // DBG("cand: %lld cost: %lld\\n", cand, cost);
      if (K / cost - 1 > 0) {
        chmax(ans, cand);
      }
      if (cand >= 0) {
        dfs2(i, cand, cost + (K % cost));
      }
    }
    // 2 周だけ
    dfs2(i, 0, min(K, 2 * cost));
  }
  return ans;
}

int bf_ans = -LLINF;
void bf_dfs(int i, int cur, int rem) {
  if (rem <= 0) {
    return;
  }
  cur += C[P[i]];
  chmax(bf_ans, cur);
  bf_dfs(P[i], cur, rem - 1);
}

int bf() {
  REP(i, N) { chmax(bf_ans, C[i]); }

  REP(i, N) { bf_dfs(i, 0, K); }

  return bf_ans;
}

signed main() {
  IOS;

#ifdef SUBMIT
  cin >> N >> K;
  P.resize(N);
  C.resize(N);
  REP(i, N) {
    cin >> P[i];
    P[i]--;
  }
  REP(i, N) { cin >> C[i]; }
  int res = solve();
  print(res);
#else

  random_device r;
  mt19937 mt(r());
  uniform_int_distribution<> dist(0, 30);
  N = 6;
  REP(loop, 10000) {
    K = dist(mt);
    // P の作成
    P.clear();
    REP(i, N) { P.push_back(i); }
    do {
      C.clear();
      REP(i, N) {
        if (dist(mt) % 2 == 0) {
          C.push_back(-1 * dist(mt));
        } else {
          C.push_back(dist(mt));
        }
      }
      int cand = solve();
      int bf_ans = bf();
      if (cand != bf_ans) {
        DBG("cand: %lld bf_ans: %lld\\n", cand, bf_ans);
        DBG("K: %lld\\n", K);
        DBG("P: ");
        REP(i, N) { DBG("%lld ", P[i]); }
        DBG("\\n");
        DBG("C: ");
        REP(i, N) { DBG("%lld ", C[i]); }
        DBG("\\n");
        return 0;
      }
    } while (next_permutation(P.begin(), P.end()));
  }
#endif

  return 0;
}


Error: TLE

Fixed submission:

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

#pragma region Macros
using ll = long long;
#define int ll
using pii = pair<int, int>;
using tiii = tuple<int, int, int>;
template <class T = ll> using V = vector<T>;
template <class T = ll> using VV = V<V<T>>;
#define IOS                                                                    \\
  ios::sync_with_stdio(false);                                                 \\
  cin.tie(0);                                                                  \\
  cout.tie(0);
#define FOR(i, l, r) for (ll i = (l); i < (r); ++i)
#define REP(i, n) FOR(i, 0, n)
#define REPS(i, n) FOR(i, 1, n + 1)
#define RFOR(i, l, r) for (ll i = (l); i >= (r); --i)
#define RREP(i, n) RFOR(i, n - 1, 0)
#define RREPS(i, n) RFOR(i, n, 1)
#define mp make_pair
#define mt make_tuple
#define pb push_back
#define eb emplace_back
#define all(x) (x).begin(), (x).end()
#define SORT(name) sort(name.begin(), name.end())
#define RSORT(name)                                                            \\
  SORT(name);                                                                  \\
  reverse(all(name));
#define ZERO(p) memset(p, 0, sizeof(p))
#define MINUS(p) memset(p, -1, sizeof(p))
inline void Yes(bool b = true) { cout << (b ? "Yes" : "No") << '\\n'; }
inline void YES(bool b = true) { cout << (b ? "YES" : "NO") << '\\n'; }
template <class T> inline void print(T x) { cout << x << '\\n'; }
template <typename T1, typename T2> inline void chmin(T1 &a, T2 b) {
  if (a > b)
    a = b;
}
template <typename T1, typename T2> inline void chmax(T1 &a, T2 b) {
  if (a < b)
    a = b;
}

const ll LLINF = (1LL << 60);
const int INF = (1LL << 30);
const double DINF = std::numeric_limits<double>::infinity();
#pragma endregion

#if 1
#define DBG(fmt, ...) printf(fmt, ##__VA_ARGS__)
#else
#define DBG(fmt, ...)
#endif

class UnionFind {
public:
  UnionFind() {}
  UnionFind(int size) { Init(size); }
  // size: Union-Find で扱いたい要素の数
  void Init(int size) { m_Data.assign(size, -1); }

  // x と y の属する集合を併合
  // 既に一緒な集合内にいるなら false を
  // 新たに併合されたなら true を返す
  bool Unite(int x, int y) {
    x = Find(x), y = Find(y);
    if (x == y) {
      return (false);
    }
    if (m_Data[x] > m_Data[y]) {
      swap(x, y);
    }
    m_Data[x] += m_Data[y];
    m_Data[y] = x;
    return (true);
  }

  // 要素 k が属する集合の根を返す
  int Find(int k) {
    if (m_Data[k] < 0) {
      return (k);
    }
    return (m_Data[k] = Find(m_Data[k]));
  }

  // 要素 k と同じ集合の要素数を返す
  int Size(int k) { return (-m_Data[Find(k)]); }

  // 要素のうち独立しているグループの数 O(NlogN)
  int GetGroupNum() {
    map<int, bool> groups;
    for (int i = 0; i < m_Data.size(); ++i) {
      groups[Find(i)] = true;
    }
    return static_cast<int>(groups.size());
  }

private:
  vector<int> m_Data;
};

#define SUBMIT

const int MOD = 1000000007;
const int MAX_N = 100010;

int N, K;
V<> P;
V<> C;

UnionFind uf;
V<bool> isUsed;
int ans = -LLINF;
V<> group; // グループ i の1周の点数

int dfs(int i) {
  if (isUsed[i]) {
    return 0;
  }
  isUsed[i] = true;

  uf.Unite(i, P[i]);
  return dfs(P[i]) + C[P[i]];
}

void dfs2(int i, int cur, int rem) {
  if (isUsed[i]) {
    return;
  }
  if (rem <= 0) {
    return;
  }
  isUsed[i] = true;

  cur += C[P[i]];
  chmax(ans, cur);
  dfs2(P[i], cur, rem - 1);
}

int solve() {
  REP(i, N) { chmax(ans, C[i]); }

  uf.Init(N);

  // group の作成
  group.assign(N, 0);
  isUsed.assign(N, false);
  REP(i, N) {
    if (isUsed[i]) {
      continue;
    }
    int val = dfs(i);
    int g_id = uf.Find(i);
    group[g_id] = val;
  }

  REP(i, N) {
    int g_id = uf.Find(i);
    // 1 周にかかるコスト
    int cost = uf.Size(i);
    isUsed.assign(N, false);
    if (group[g_id] > 0) {
      // ぐるぐる回る
      int cand = K / cost * group[g_id];
      chmax(ans, cand);
      dfs2(i, cand, K % cost);
    } else {
      // 1 周だけ
      dfs2(i, 0, min(K, cost));
    }
  }
  return ans;
}

int bf_ans = -LLINF;
void bf_dfs(int i, int cur, int rem) {
  if (rem <= 0) {
    return;
  }
  cur += C[P[i]];
  chmax(bf_ans, cur);
  bf_dfs(P[i], cur, rem - 1);
}

int bf() {
  REP(i, N) { chmax(bf_ans, C[i]); }

  REP(i, N) { bf_dfs(i, 0, K); }

  return bf_ans;
}

signed main() {
  IOS;

#ifdef SUBMIT
  cin >> N >> K;
  P.resize(N);
  C.resize(N);
  REP(i, N) {
    cin >> P[i];
    P[i]--;
  }
  REP(i, N) { cin >> C[i]; }
  // int res = solve();
  int res = bf();
  print(res);
#else
#endif

  return 0;
}