AtCoder Regular Contest 018 C: 席替え

問題

C: 席替え - AtCoder Regular Contest 018 | AtCoder

解法

重要な制約として"p <=N*M" というのがあり、a%p==0 で場合分けしなければならない（ここでめちゃくちゃハマった）。

a%p!=0 の時、すべての生徒の成績は異なるため、各々が前から何行目に来るかが定まる。各生徒を自分がいるべき行まで移動させる。この時、列は移動させないため、重なる生徒が出る。各行について、重なっている生徒を一番近くの相手いる席まで移動させてやれば良い。

コード

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
template <typename T>
std::ostream &operator<<(std::ostream &out, const std::vector<T> &v) {
  if (!v.empty()) {
    out << '[';
    std::copy(v.begin(), v.end(), std::ostream_iterator<T>(out, ", "));
    out << "\b\b]";
  }
  return out;
}
template <typename T1, typename T2>
std::ostream &operator<<(std::ostream &out, const std::pair<T1, T2> &p) {
  out << "[" << p.first << ", " << p.second << "]";
  return out;
}
template <class T, class U>
void chmin(T &t, U f) {
  if (t > f) t = f;
}
template <class T, class U>
void chmax(T &t, U f) {
  if (t < f) t = f;
}
int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  ll N, M, x, a, p;
  cin >> N >> M >> x >> a >> p;
  if (a % p == 0) {
    if (x >= p) {
      cout << 2 * (N - 1) << endl;
    } else {
      cout << 0 << endl;
    }
    return 0;
  }

  vector<ll> score(N * M);
  score[0] = x;
  for (int i = 1; i < N * M; ++i) score[i] = (score[i - 1] + a) % p;
  vector<ll> cx = score;
  sort(cx.begin(), cx.end());
  map<ll, ll> rank;
  for (int i = 0; i < N * M; ++i) rank[cx[i]] = i / M;
  for (int i = 0; i < N * M; ++i) score[i] = rank[score[i]];

  ll ans = 0;
  vector<vector<ll>> rs(N);
  for (int m = 0; m < M; ++m) {
    for (int r = 0; r < N; ++r) {
      int s = score[r * M + m];
      ans += abs(s - r);
      assert(s < N);
      rs[s].push_back(m);
    }
  }

  for (int r = 0; r < N; ++r) {
    sort(rs[r].begin(), rs[r].end());
    assert(rs[r].size() == M);
    for (int i = 0; i < M; ++i) {
      ans += abs(i - rs[r][i]);
    }
  }
  cout << ans << endl;
}