Pythonで二値グリッドを整列させるための最小スワップ回数を求めるプログラム

問題の概要

n × n の二値(0と1のみ)行列を考えます。この行列に対して、「隣接する2つの行を選んで入れ替える」という操作を1ステップとして実行できます。ここで求めたいのは、行列の主対角線より上側にあるすべての要素が 0 になるようにするために必要な最小スワップ回数です。どのように行を入れ替えても条件を満たせない場合は、-1 を返します。

たとえば、次のような入力が与えられたとします。

この場合、出力は 2 になります。2回の隣接スワップで行を並べ替えれば、主対角線より上の要素をすべて 0 にできるからです。

解き方のポイント

この問題を効率よく解く鍵は、各行を「右端にいくつ 0 が連続しているか」という情報に変換することです。具体的には、各行について最も左にある 1 より右側に続く要素の個数(= 後続の 0 の数)をあらかじめ計算しておきます。

i 行目(0始まり)に置かれた行が条件を満たすには、列 i+1 以降がすべて 0 である必要があります。つまり、後続の 0 の数が「n − i − 1」以上であればよいことになります。この性質を使い、上から順に行を確定させていく貪欲法(グリーディー法)で最小スワップ数を求めます。

アルゴリズムの手順

• n := 行列の行数とします。
• m := サイズ n の配列を作成し、すべて n で初期化します。
• i を 0 から n−1 まで繰り返します:
  • j を n−1 から 0 まで減らしながら繰り返します:
    • matrix[i, j] が 1 ならば:
      • m[i] := n − j − 1 とし、内側のループを抜けます。
• t := 0、ans := 0 とします。
• i を 0 から n−1 まで繰り返します:
  • t := t + 1 とします。
  • flag := False とします。
  • j を i から n−1 まで繰り返します:
    • m[j] ≥ n − t ならば:
      • ans := ans + (j − i) とします(この行を i の位置まで移動させるのに必要なスワップ回数)。
      • flag := True とし、ループを抜けます。
  • flag が False のままなら、条件を満たす行が存在しないため -1 を返します。
  • m のインデックス i+1〜j を、m のインデックス i〜j−1 の内容で更新します(行を実際に移動させた状態を再現)。
• 最後に ans を返します。

実装例

それでは、理解を深めるために実際のコードを見てみましょう。

def solve(matrix):
    n = len(matrix)
    m = [n] * n
    for i in range(n):
        for j in range(n-1,-1,-1):
            if matrix[i][j] == 1:
                m[i] = n-j-1
                break
    t,ans = 0,0
    for i in range(n):
        t += 1
        flag = False
        for j in range(i,n):
            if m[j] >= n-t:
                ans += j-i
                flag = True
                break
        if not flag: return -1
        m[i+1:j+1] = m[i:j]
    return ans
matrix = [[0,1,0],[0,1,1],[1,0,0]]
print(solve(matrix))

入力

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

出力

2

計算量について

前処理で各行の後続ゼロ数を求める部分は O(n²)、貪欲にマッチングを行う部分も O(n²) なので、全体の時間計算量は O(n²) となります。空間計算量は O(n) です。隣接スワップの性質上、ある行を k 行上へ移動させるには必ず k 回のスワップが必要なため、この貪欲な選び方が最小スワップ数を保証します。