C++で行と列を反転して2値行列の「1」の数を最大化するアルゴリズム

2値(0と1のみで構成される)行列が与えられたとき、1つの行を反転し、その後1つの列を反転することで、得られる「1」の最大数を求める問題をC++で解く方法を解説します。

問題の例

たとえば、次のような入力行列を考えてみましょう。

この場合、最適な行と列を選んで反転すると、出力は 8 になります。

解法のアプローチ

この問題は、全探索を行う代わりに、行ごと・列ごとの「1」の個数を事前に集計することで効率的に解けます。手順は以下の通りです。

• n := 行列の行数

• m := 行列の列数

• ret := 0(答えを格納する変数)

• サイズ n の配列 row を定義(各行に含まれる「1」の個数)

• サイズ m の配列 col を定義(各列に含まれる「1」の個数)

• total := 0(行列全体の「1」の総数)

• i := 0 から n 未満の間、以下を繰り返す:

  • j := 0 から m 未満の間、以下を繰り返す:

    • row[i] := row[i] + matrix[i, j]

    • col[j] := col[j] + matrix[i, j]

    • total := total + matrix[i, j]

• 再び i := 0 から n 未満の間、以下を繰り返す:

  • j := 0 から m 未満の間、以下を繰り返す:

    • cand := total - row[i] - col[j] + ((m - row[i]) + (n - col[j]))

    • matrix[i, j] が 0 以外の場合:

      • cand := cand + 2

    • それ以外の場合:

      • cand := cand - 2

    • ret := ret と cand のうち大きい方

• ret を返す

考え方のポイント

行 i を反転すると、その行に含まれる row[i] 個の「1」はすべて「0」に変わり、残りの (m − row[i]) 個の「0」はすべて「1」に変わります。同様に、列 j を反転すると、col[j] 個の「1」が「0」に、(n − col[j]) 個の「0」が「1」に変わります。

ただし、セル (i, j) は行の反転と列の反転の両方の影響を受けるため、上記の計算では二重に扱われてしまいます。そこで、matrix[i][j] が 1 の場合は cand に 2 を加え、0 の場合は 2 を引いて補正します。これにより、すべての (i, j) の組み合わせについて反転後の「1」の数を正しく求められます。

計算量は前処理・探索ともに O(n × m) となり、非常に効率的です。

実装例

理解を深めるために、以下のC++の実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
   public:
   int solve(vector<vector<int>> &matrix) {
      int n = matrix.size();
      int m = matrix[0].size();
      int ret = 0;
      vector<int> row(n);
      vector<int> col(m);
      int total = 0;
      for (int i = 0; i < n; i++) {
         for (int j = 0; j < m; j++) {
            row[i] += matrix[i][j];
            col[j] += matrix[i][j];
            total += matrix[i][j];
         }
      }
      for (int i = 0; i < n; i++) {
         for (int j = 0; j < m; j++) {
            int cand = total - row[i] - col[j] + (m - row[i]) + (n -
            col[j]);
            if (matrix[i][j]) {
               cand += 2;
            }else {
               cand -= 2;
            }
            ret = max(ret, cand);
         }
      }
      return ret;
   }
};
main() {
   Solution ob;
   vector<vector<int>> v = {{1,0,1},{0,1,0},{1,0,0}};
   cout << (ob.solve(v));
}

入力

{{1,0,1},{0,1,0},{1,0,0}}

出力

8