Problem p02964 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <set>
#include <tuple>
#include <vector>

using namespace std;

#define rep(i, n) for (int64_t i = 0; i < (n); i++)
#define irep(i, n) for (int64_t i = 0; i <= (n); i++)
#define rrep(i, n) for (int64_t i = (n)-1; i >= 0; i--)
#define rirep(i, n) for (int64_t i = n; i >= 0; i--)

struct Permutation {
  vector<int64_t> container;
  Permutation(size_t size) : container(size) {
    for (int i = 0; i < size; i++) {
      container[i] = 0;
    }
  }
  size_t size() const { return container.size(); }
  int64_t &operator[](size_t i) { return container[i]; }
  const int64_t &operator[](size_t i) const { return container[i]; }
  Permutation operator*(const Permutation &other) const {
    int64_t n = container.size();
    Permutation result(n);
    for (int64_t i = 0; i < n; i++) {
      result[i] = (*this)[i] + other[(i + (*this)[i]) % n];
    }
    return result;
  }
};

Permutation ppow(const Permutation &p, int64_t x) {
  if (x == 0) {
    return Permutation(p.size());
  } else if (x % 2 == 0) {
    return ppow(p * p, x / 2);
  } else {
    return p * ppow(p, x - 1);
  }
}

class UFTree {
  vector<int> parent;
  vector<int> height;

public:
  UFTree(int size) {
    for (int i = 0; i < size; i++) {
      parent.push_back(i);
      height.push_back(0);
    }
  }
  int root(int node) {
    if (parent[node] == node) {
      return node;
    }
    return parent[node] = root(parent[node]);
  }
  void merge(int n0, int n1) {
    int r0 = root(n0), r1 = root(n1);
    if (r0 == r1) {
      return;
    }
    if (height[r0] < height[r1])
      swap(r0, r1);
    if (height[r0] == height[r1])
      height[r0]++;
    parent[r1] = r0;
  }
};

int main() {
  int n;
  int64_t k;
  cin >> n >> k;

  vector<int> a(n);
  rep(i, n) {
    cin >> a[i];
    a[i]--;
  }

  vector<vector<int>> idx(200'000);
  rep(i, n) { idx[a[i]].push_back(i); }

  Permutation p(n);
  rep(i, idx.size()) {
    int m = idx[i].size();
    rep(j, m) { p[idx[i][j]] = (idx[i][(j + 1) % m] - idx[i][j] + 1 + n) % n; }
  }
  rep(i, n) {
    if (p[i] == 1) {
      p[i] = n + 1;
    }
  }

  UFTree uft(n);
  rep(i, n) { uft.merge(i, (i + p[i]) % n); }
  int64_t T = 0;
  int cnt = 0;
  rep(i, n) {
    if (uft.root(0) == uft.root(i)) {
      T += p[i];
      cnt++;
    }
  }

  int64_t M = (k * n) % T;
  int64_t lb = 0, ub = cnt;
  while (ub > lb + 1) {
    int64_t mid = (lb + ub) / 2;

    Permutation q = ppow(p, mid);
    if (q[0] <= M) {
      lb = mid;
    } else {
      ub = mid;
    }
  }
  Permutation q = ppow(p, lb);

  vector<int> result;
  int64_t j = q[0];
  while (j < M) {
    if (j + p[j % n] <= M) {
      j += p[j % n];
    } else {
      result.push_back(a[j % n]);
      j++;
    }
  }

  rep(i, result.size()) { cout << result[i] + 1 << " "; }
  cout << endl;

  return 0;
}

Fixed submission:
#include <algorithm>
#include <cmath>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <set>
#include <tuple>
#include <vector>

using namespace std;

#define rep(i, n) for (int64_t i = 0; i < (n); i++)
#define irep(i, n) for (int64_t i = 0; i <= (n); i++)
#define rrep(i, n) for (int64_t i = (n)-1; i >= 0; i--)
#define rirep(i, n) for (int64_t i = n; i >= 0; i--)

struct Permutation {
  vector<int64_t> container;
  Permutation(size_t size) : container(size) {
    for (int i = 0; i < size; i++) {
      container[i] = 0;
    }
  }
  size_t size() const { return container.size(); }
  int64_t &operator[](size_t i) { return container[i]; }
  const int64_t &operator[](size_t i) const { return container[i]; }
  Permutation operator*(const Permutation &other) const {
    int64_t n = container.size();
    Permutation result(n);
    for (int64_t i = 0; i < n; i++) {
      result[i] = (*this)[i] + other[(i + (*this)[i]) % n];
    }
    return result;
  }
};

Permutation ppow(const Permutation &p, int64_t x) {
  if (x == 0) {
    return Permutation(p.size());
  } else if (x % 2 == 0) {
    return ppow(p * p, x / 2);
  } else {
    return p * ppow(p, x - 1);
  }
}

class UFTree {
  vector<int> parent;
  vector<int> height;

public:
  UFTree(int size) {
    for (int i = 0; i < size; i++) {
      parent.push_back(i);
      height.push_back(0);
    }
  }
  int root(int node) {
    if (parent[node] == node) {
      return node;
    }
    return parent[node] = root(parent[node]);
  }
  void merge(int n0, int n1) {
    int r0 = root(n0), r1 = root(n1);
    if (r0 == r1) {
      return;
    }
    if (height[r0] < height[r1])
      swap(r0, r1);
    if (height[r0] == height[r1])
      height[r0]++;
    parent[r1] = r0;
  }
};

int main() {
  int n;
  int64_t k;
  cin >> n >> k;

  vector<int> a(n);
  rep(i, n) {
    cin >> a[i];
    a[i]--;
  }

  if (n == 1) {
    if (k % 2 == 0) {
      cout << endl;
    } else {
      cout << a[0] + 1 << endl;
    }
    return 0;
  }

  vector<vector<int>> idx(200'000);
  rep(i, n) { idx[a[i]].push_back(i); }

  Permutation p(n);
  rep(i, idx.size()) {
    int m = idx[i].size();
    rep(j, m) { p[idx[i][j]] = (idx[i][(j + 1) % m] - idx[i][j] + 1 + n) % n; }
  }
  rep(i, n) {
    if (p[i] == 1) {
      p[i] = n + 1;
    }
  }

  UFTree uft(n);
  rep(i, n) { uft.merge(i, (i + p[i]) % n); }
  int64_t T = 0;
  int cnt = 0;
  rep(i, n) {
    if (uft.root(0) == uft.root(i)) {
      T += p[i];
      cnt++;
    }
  }

  int64_t M = (k * n) % T;
  int64_t lb = 0, ub = cnt;
  while (ub > lb + 1) {
    int64_t mid = (lb + ub) / 2;

    Permutation q = ppow(p, mid);
    if (q[0] <= M) {
      lb = mid;
    } else {
      ub = mid;
    }
  }
  Permutation q = ppow(p, lb);

  vector<int> result;
  int64_t j = q[0];
  while (j < M) {
    if (j + p[j % n] <= M) {
      j += p[j % n];
    } else {
      result.push_back(a[j % n]);
      j++;
    }
  }

  rep(i, result.size()) { cout << result[i] + 1 << " "; }
  cout << endl;

  return 0;
}