Pythonで2値行列内の「特別な位置」の数を求めるプログラム

問題の概要

m × n の2値(0と1のみで構成された)行列が与えられたとき、その中に含まれる「特別な位置」の総数を求めることを考えます。ここで、位置 (i, j) が「特別」とみなされるのは、mat[i][j] = 1 であり、かつ i 行目と j 列目にあるその他のすべての要素が 0 である場合です。つまり、その行と列の中で唯一の 1 になっているマスを探す問題です。

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

この場合の出力は 3 になります。特別な位置は (0, 0)、(1, 2)、(3, 1) の3箇所です。

解き方のアルゴリズム

この問題は、次の手順で解くことができます。

• カウンター special を 0 で初期化します。
• 行列の各行 i に対して、次の処理を繰り返します。
  • 行 matrix[i] に含まれる 1 の個数がちょうど 1 であるかを確認します。
  • 1 が 1 つだけの場合、その位置 indexOfOne を取得し、同じ列に 1 がいくつ存在するかを数えます。
  • 列の走査中に 1 が 2 個以上見つかった時点で、その位置は特別ではないため、ループを抜けて無駄な計算を省きます。
  • 列に含まれる 1 がちょうど 1 個だけだった場合、位置 (i, indexOfOne) は特別な位置なので、special を 1 増やします。
• すべての行を調べ終えたら、special の値を返します。

Pythonでの実装例

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

def solve(matrix):
    special = 0
    for i in range(len(matrix)):
        if matrix[i].count(1) == 1:
            numOfOne = 0
            indexOfOne = matrix[i].index(1)
            for j in range(len(matrix)):
                if matrix[j][indexOfOne] == 1:
                    numOfOne += 1
                if numOfOne > 1:
                    break

            if numOfOne == 1:
                special += 1

    return special

matrix = [[1,0,0,0,0],
          [0,0,1,0,0],
          [0,0,0,1,1],
          [0,1,0,0,0]]
print(solve(matrix))

実行結果

入力:

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

出力:

3

ポイントのまとめ

このアルゴリズムでは、まず「行に 1 が 1 つしかないこと」で候補を絞り込み、その後で「列にも 1 が 1 つしかないこと」を確認するという2段階のチェックを行っています。条件を満たさないことが確定した時点で break して走査を打ち切ることで、無駄な計算を抑えられる点が工夫されています。また、Python の list.count() と list.index() を活用することで、ロジックを簡潔かつ読みやすく記述できるのも大きな魅力です。