C++で行列のすべての要素を等しくするために必要な最小演算回数を求める

問題概要

整数 K と M × N の行列が与えられます。1回の演算では、行列内の任意の要素に対して K を足すか引くことができます。このとき、行列のすべての要素を等しい値に揃えるために必要な最小の演算回数を求めるのが課題です。

例

たとえば、入力行列と K が次のとおりである場合を考えます。

入力行列:
{
    {2, 4},
    {20, 40}
}
K = 2

この場合、すべての要素を中央値の 20 に揃えることで、合計 27 回の演算で実現できます。

Matrix[0][0]: 2 + (K × 9)   = 20 → 9 回
Matrix[0][1]: 4 + (K × 8)   = 20 → 8 回
Matrix[1][1]: 40 − (K × 10) = 20 → 10 回

合計: 9 + 8 + 10 = 27 回

アルゴリズム

最適解を効率よく求めるには、次の手順に従います。

1. 剰余の一致を確認する: 各要素に加減できるのは K の倍数だけなので、すべての要素を K で割った余り（mod）は互いに等しくなければなりません。等しくない場合は、どれだけ演算を繰り返しても全要素を一致させることができないため、-1 を返します。
2. 中央値を求める: 行列の全要素を1次元配列に展開し、昇順（非降順）にソートしたうえで中央値を取得します。絶対差の総和を最小にする値が中央値になるという性質を利用しています。
3. 全要素を中央値へ揃える: 各要素について必要な演算回数は |要素 − 中央値| ÷ K であり、その総和が答えとなります。

C++実装例

#include <bits/stdc++.h>
using namespace std;

int getMinOperations(int n, int m, int k, vector<vector<int>>& matrix) {
    vector<int> arr(n * m, 0);
    int mod = matrix[0][0] % k;

    // すべての要素の K による余りが一致するかを確認
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < m; ++j) {
            arr[i * m + j] = matrix[i][j];
            if (matrix[i][j] % k != mod) {
                return -1;
            }
        }
    }

    // 昇順にソートして中央値を求める
    sort(arr.begin(), arr.end());
    int median = arr[(n * m) / 2];

    int minOperations = 0;
    for (int i = 0; i < n * m; ++i)
        minOperations += abs(arr[i] - median) / k;

    // 要素数が偶数の場合は、もう1つの中央値候補も比較する
    if ((n * m) % 2 == 0) {
        int newMedian = arr[(n * m) / 2 - 1];
        int newMinOperations = 0;
        for (int i = 0; i < n * m; ++i)
            newMinOperations += abs(arr[i] - newMedian) / k;
        minOperations = min(minOperations, newMinOperations);
    }
    return minOperations;
}

int main() {
    vector<vector<int>> matrix = {
        {2, 4},
        {20, 40},
    };
    int n = matrix.size();
    int m = matrix[0].size();
    int k = 2;
    cout << "最小演算回数 = " << getMinOperations(n, m, k, matrix) << endl;
    return 0;
}

このプログラムをコンパイルして実行すると、次の出力が得られます。

出力

最小演算回数 = 27

計算量

処理の中心はソートであるため、時間計算量は O(N·M log(N·M))、ソート用の一時配列が必要となるため空間計算量は O(N·M) です。