Pythonで正方形の部分行列を転置して別の行列へ変換できるかを判定する方法

N × M の2つの行列 mat1 と mat2 が与えられているとします。許されている操作は「mat1 内の任意の正方形の部分行列を選び、それを転置する」というものです。この操作を何度でも適用できるとき、mat1 から mat2 を作り出せるかどうかを判定するのがこの問題です。

たとえば、入力が次のようなケースを考えてみましょう。

この場合の出力は True になります。mat1 の右上側にある 2×2 の部分行列を転置すると、mat2 と完全に一致するからです。

鍵となる性質：「行番号＋列番号」は転置しても変わらない

まず押さえておきたい重要な観察点があります。位置 (i, j) の要素を転置すると、その要素は (j, i) へ移動します。このとき「行番号＋列番号」、つまり i + j の値は一切変わりません。

正方形の部分行列内での転置でも事情は同じで、要素は必ず同じ反対角線（右肩上がりの斜めライン）上にとどまります。言い換えると、

• 同じ反対角線上の要素同士であれば、転置を組み合わせることで自由に入れ替えられる
• 異なる反対角線の間で要素を移動させることは絶対にできない

したがって、「mat1 を mat2 に変換できる」ための条件は、次のようにシンプルに言い換えられます。

すべての反対角線について、mat1 と mat2 に含まれる要素の多重集合が一致していること

解法のステップ

上記の性質を踏まえると、アルゴリズムは次の手順で組み立てられます。

• row := 行列の行数、column := 行列の列数
• i を 0 ～ row - 1 の範囲で繰り返す（左端の列を起点とする反対角線を処理）
  • temp1 := 新しいリスト、temp2 := 新しいリスト
  • r := i、col := 0 とする
  • r >= 0 かつ col < column の間、次を繰り返す
    • temp1 に mat1[r][col] を追加
    • temp2 に mat2[r][col] を追加
    • r := r - 1、col := col + 1
  • temp1 と temp2 をそれぞれソートする
  • i を 0 ～ temp1 のサイズ - 1 の範囲で繰り返し、temp1[i] が temp2[i] と異なる場合は False を返す
• j を 1 ～ column - 1 の範囲で繰り返す（下端の行を起点とする残りの反対角線を処理）
  • r := row - 1、col := j から同様に右上方向へ要素を収集し、ソート後に一致しなければ False を返す
• すべての反対角線が一致していれば True を返す

ソートしてから比較しているのは、転置によって反対角線上の要素の並び順が変わっても構わないためです。順序を無視した「要素の中身」だけを揃えて比較することで、判定漏れや誤判定を防いでいます。

実装例

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

def solve(mat1, mat2):
    row = len(mat1)
    column = len(mat1[0])
    # 左端の列を起点とする反対角線をチェック
    for i in range(row):
        temp1 = []
        temp2 = []
        r = i
        col = 0
        while r >= 0 and col < column:
            temp1.append(mat1[r][col])
            temp2.append(mat2[r][col])
            r -= 1
            col += 1
        temp1.sort()
        temp2.sort()
        for i in range(len(temp1)):
            if temp1[i] != temp2[i]:
                return False
    # 下端の行を起点とする残りの反対角線をチェック
    for j in range(1, column):
        temp1 = []
        temp2 = []
        r = row - 1
        col = j
        while r >= 0 and col < column:
            temp1.append(mat1[r][col])
            temp2.append(mat2[r][col])
            r -= 1
            col += 1
        temp1.sort()
        temp2.sort()
        for i in range(len(temp1)):
            if temp1[i] != temp2[i]:
                return False
    return True

mat1 = [
    [5, 6, 7],
    [1, 2, 3],
    [6, 8, 9]]
mat2 = [
    [5, 6, 2],
    [1, 7, 3],
    [6, 8, 9]]
print(solve(mat1, mat2))

入力

[
    [5, 6, 7],
    [1, 2, 3],
    [6, 8, 9]],
[
    [5, 6, 2],
    [1, 7, 3],
    [6, 8, 9]]

出力

True

計算量について

すべての要素は各反対角線でちょうど1回ずつ走査され、反対角線ごとにソートが行われます。反対角線の長さは高々 min(N, M) であるため、全体の計算量は O(N・M・log(min(N, M))) 程度に収まり、行列サイズがそこまで大きくなければ十分高速に動作します。

空間計算量についても、一時的に保持するのは1本分の反対角線の要素だけなので、O(min(N, M)) で済む点も嬉しいところです。