Pythonで行と列の合計を満たす有効な行列を見つけるプログラム

問題の概要

2つの配列 rowSum と colSum があり、それぞれ非負の整数が格納されているとします。rowSum[i] は2次元行列の i 行目の要素の合計を、colSum[j] は j 列目の要素の合計を表します。

このとき、与えられた rowSum と colSum の条件をすべて満たすような、非負の値のみで構成されたサイズ（rowSumの長さ × colSumの長さ）の行列を1つ見つける必要があります。

例として、入力が rowSum = [13,14,12]、colSum = [9,13,17] の場合、出力は次のようになります。

この行列では、各行の合計が 13・14・12 に、各列の合計が 9・13・17 になっており、条件を満たしていることが確認できます。

解法のアプローチ（貪欲法）

この問題は貪欲法（グリーディー法）で解くことができます。基本的な考え方は、「まだ確定していない行・列の中で最も合計値が小さいものを選び、その値を未処理の相手側にできるだけ割り当てる」というものです。これにより、必ず条件を満たす行列を構築できます。

アルゴリズムの手順

• matrix := 全要素が0の空の行列を作成する
• visited := 処理済みの行・列を記録する新しいセット
• minimum(r, c) 関数を定義する
• min_total := 無限大（∞）で初期化
• type := 空文字列で初期化
• i を 0 から r の長さ - 1 まで繰り返す：
  • r[i] < min_total ならば
    • index := i、type := 'row'、min_total := r[i]
• i を 0 から c の長さ - 1 まで繰り返す：
  • c[i] < min_total ならば
    • min_total := c[i]、type := 'col'、index := i
• type が 'row' の場合：
  • r[index] := 無限大に設定（処理済みマーク）
  • i を 0 から c の長さ - 1 まで繰り返し、c[i] が無限大ではなく c[i] >= min_total である最初の位置で、c[i] から min_total を減算し、matrix[index][i] に min_total を代入してループを抜ける
• type が 'col' の場合：
  • c[index] := 無限大に設定（処理済みマーク）
  • i を 0 から r の長さ - 1 まで繰り返し、r[i] が無限大ではなく r[i] >= min_total である最初の位置で、r[i] から min_total を減算し、matrix[i][index] に min_total を代入してループを抜ける
• (index, type) のペアを visited に追加する
• メイン処理では、visited のサイズが len(r) + len(c) と等しくなるまで minimum(r, c) を繰り返し呼び出す
• 最後に matrix を返す

実装例

それでは、上記のアルゴリズムをPythonで実装してみましょう。

def solve(r, c):
   matrix = [[0]*len(c) for _ in range(len(r))]
   visited = set()

   def minimum(r,c):
      min_total = float('inf')
   
      type = ''
      for i in range(len(r)):
         if(r[i] < min_total):
            index = i
            type = 'row'
            min_total = r[i]

      for i in range(len(c)):
         if(c[i] < min_total):
            min_total = c[i]
            type = 'col'
            index = i

      if(type == 'row'):
         r[index] = float('inf')

         for i in range(len(c)):
            if(c[i] != float('inf') and c[i] >= min_total):
               c[i] -= min_total
               matrix[index][i] = min_total
               break

      if(type == 'col'):
         c[index] = float('inf')
         for i in range(len(r)):
            if(r[i] != float('inf') and r[i] >= min_total):
               r[i] -= min_total
               matrix[i][index] = min_total
               break

      visited.add((index,type))

   while len(visited) != len(r)+len(c):
      minimum(r,c)

   return matrix

rowSum = [13,14,12]
colSum = [9,13,17]
print(solve(rowSum, colSum))

入力

[13,14,12], [9,13,17]

出力

[[9, 4, 0], [0, 9, 5], [0, 0, 12]]

まとめ

このアルゴリズムは、毎回残りの行・列の中で最小の合計値を持つものを優先的に処理することで、すべての制約を確実に満たす行列を構築します。計算量は O((m+n)²) 程度であり、m を行数、n を列数とした場合でも十分に高速に動作します。答えが一意ではない場合でも、条件を満たす行列のうちの1つを返せばよいため、この貪欲的な手法が有効です。