AtCoder Regular Contest 018 C - 席替え
解法
まず全体を(値, 行, 列) についてソートすることによって, 生徒がそれぞれどの行に行くか, というのが決定します(同じ点数の場合は, 行順にソートするのが明らかに得)。
で, そうすると各行に対して, 「誰をどの列に配置するか」という問題になります。これは, 列順にソートして 0, 1, ..., M に並べていけば良いです。
計算量は O(NM log (NM) + NM log (M)) なので間に合いますね。
int main() { cin.tie(0); ios::sync_with_stdio(false); int N, M; cin >> N >> M; vector<pair<int, pii> > v(N*M); int x, a, p; cin >> x >> a >> p; for (int i = 0; i < N*M; i++) { v[i] = make_pair(x, pii(i/M, i%M)); x = (x+a)%p; } sort(v.begin(), v.end()); ll ans = 0; for (int i = 0; i < N; i++) { vector<int> vv(M); for (int j = 0; j < M; j++) { auto p = v[i*M+j]; ans += abs(p.second.first - i); vv[j] = p.second.second; } sort(vv.begin(), vv.end()); for (int j = 0; j < M; j++) ans += abs(vv[j]-j); } cout << ans << endl; return 0; }