C++で2行のバイナリ行列を再構築するアルゴリズムと実装例

問題の概要

ここでは、n列×2行の行列について、次の情報が与えられているものとします。

• 行列の各要素は 0 または 1 のみ
• 0行目(上段)の要素の合計が upper として与えられる
• 1行目(下段)の要素の合計が lower として与えられる
• i列目(0始まり)の要素の合計が colsum[i] として与えられる(colsum は長さ n の整数配列)

この課題では、upper、lower、colsum の情報をもとに元の行列を復元し、2次元の整数配列として返します。有効な解が複数存在する場合は、そのうちのどれを返しても構いません。一方、条件を満たす解が存在しない場合は、空の2次元配列を返します。

たとえば、upper = 2、lower = 1、colsum = [1, 1, 1] が入力された場合、出力は [[1,1,0],[0,0,1]] となります。上段の合計が 2、下段の合計が 1、各列の合計もすべて条件どおりになっていることが確認できます。

解き方(アルゴリズム)

この問題は、各列の合計値に応じて貪欲法(グリーディ法)で上段・下段に 1 を割り当てていくことで解けます。具体的な手順は以下の通りです。

• flag := true、n := 配列 c のサイズ とし、2 × n の配列 ans を用意する。
• i を 0 から n-1 まで繰り返し処理する。
  • c[i] = 2 の場合:
    • u と l をそれぞれ 1 減らす。
    • u < 0 または l < 0 になったら flag := false とする。
    • ans[0][i] = 1、ans[1][i] = 1 を設定する。
  • c[i] = 1 の場合:
    • u > l ならば u を 1 減らして ans[0][i] := 1 とする。
    • u < l ならば l を 1 減らして ans[1][i] := 1 とする。
    • u = l の場合:
      • u > 0 ならば u を 1 減らして ans[0][i] := 1 とする。
      • そうでなく l > 0 ならば l を 1 減らして ans[1][i] := 1 とする。
      • どちらでもなければ flag := false とする。
  • c[i] = 0 の場合:
    • 両行とも 0 のまま何も設定しない。
  • それ以外の不正な値の場合は flag := false とする。
• flag が false のとき、または u ≠ 0 もしくは l ≠ 0 のときは空の配列を返す。
• それ以外は ans を返す。

ポイントは、列の合計が 2 なら必ず両方の行に 1 を置き、合計が 1 の場合は残りの必要数が多い方の行へ優先的に 1 を割り当てる点です。これにより、最後まで処理したときに upper と lower をちょうど使い切れるかどうかを判定できます。

実装例

それでは、実際のC++コードを見て理解を深めましょう。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<vector<auto>> v){
   cout << "[";
   for(int i = 0; i<v.size(); i++){
      cout << "[";
      for(int j = 0; j <v[i].size(); j++){
         cout << v[i][j] << ", ";
      }
      cout << "],";
   }
   cout << "]"<<endl;
}
class Solution {
   public:
   vector<vector<int>> reconstructMatrix(int u, int l, vector<int>& c) {
      bool flag = true;
      int n = c.size();
      vector<vector<int>> ans(2, vector<int>(n));
      for(int i = 0; i < n; i++){
         if(c[i] == 2){
            u--;
            l--;
            if(u<0 || l<0)flag = false;
            ans[0][i] = 1;
            ans[1][i] = 1;
         }else if(c[i] == 1){
            if(u>l){
               u--;
               ans[0][i] = 1;
            }else if(u<l){
               l--;
               ans[1][i] = 1;
            }else{
               if(u>0){
                  u--;
                  ans[0][i] = 1;
               }else if(l > 0){
                  l--;
                  ans[1][i] = 1;
               }else
                  flag = false;
            }
         }else if(c[i] == 0){
            // 列の合計が0のため、両行とも0のまま何もしない
         }else{
            flag = false;
         }
      }
      if(!flag || u!=0 ||l!=0 )return {};
      return ans;
   }
};
main(){
   vector<int> v = {1,1,1};
   Solution ob;
   print_vector(ob.reconstructMatrix(2,1,v));
}

入力

2
1
[1,1,1]

出力

[[1, 1, 0],[0, 0, 1]]

このように、貪欲法を用いることで O(n) の計算量で行列を復元でき、解が存在しない場合も適切に空配列を返せる実装になっています。