AtCoder Regular Contest 018 C - 席替え

問題

arc018.contest.atcoder.jp

解法

まず全体を(値, 行, 列) についてソートすることによって, 生徒がそれぞれどの行に行くか, というのが決定します(同じ点数の場合は, 行順にソートするのが明らかに得)。

で, そうすると各行に対して, 「誰をどの列に配置するか」という問題になります。これは, 列順にソートして 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;
}