ARC018 [C]席替え

ARC018 の C 問題「席替え」を（解説を見ながら）解きました。
i 行 j 列 (i,j) の人の成績 G(i, j) とするとき、 (a,b) < (c,d) ならば G(a,b) < G(c,d) となるように席替えを並び替えるための最小のコストを求める問題。
まず、 x_i = (x_(i-1) + a) % p で、 a%p!=0, m*n<=p, p is prime を満たすとき、x_i(0<=i<n*m) は全て異なる値をとる、ということに理解はできるが気がつかない。数学の基礎知識が身についていない証拠。
成績がすべて異なる値をとるとき、行は floor(G[i][j]/m) で確定し、列は交差しないように小さい順にソートする。
つまり成績順にソートしたあと、各行で列順にソートしてマンハッタン距離を求めれば良い。

algorithm のソートを初めて使いました。
sort(ソート始めの位置, ソート終わりの位置+1, 比較関数);
比較するための関数を自作して指定することができるので、かなり使いやすいと思います。

#include <iostream>
#include <cstdlib>
#include <algorithm>
using namespace std;

struct Data{
    int grade;
    int i, j;

    Data(){
    }
};

bool grade_lt(const Data &l, const Data &r){
    return l.grade < r.grade;
}

bool column_lt(const Data &l, const Data &r){
    return l.j < r.j;
}

int main(){
    int n, m, x0, a, p;
    Data *datas;
    int cost;

    cin >> n >> m >> x0 >> a >> p;

    if(a % p == 0){
        cout << (x0 >= p ? 2 * n - 2 : 0) << endl;
        return 0;
    }

    datas = new Data[n * m];
    datas[0].grade = x0;
    datas[0].i = 0;
    datas[0].j = 0;

    for(int i = 1; i < n * m; i++){
        x0 += a;
        x0 %= p;
        datas[i].grade = x0;
        datas[i].i = i / m;
        datas[i].j = i % m;
    }

    sort(datas, datas+n*m, grade_lt);

    for(int i = 0; i < n; i++){
        sort(datas+i*m, datas+(i+1)*m, column_lt);

        for(int j = 0; j < m; j++){
            cost += abs(i - datas[i * m + j].i) + abs(j - datas[i * m + j].j);
        }
    }

    cout << cost << endl;

    delete [] datas;

    return 0;
}