Pythonで受け入れられる招待状の数を求めるプログラム ― DFSによる二部グラフマッチング

パーティーの準備で、m人の男子とn人の女子がいるとします(m = n)。各男子は必ず女子を一人連れて参加しなければならず、男子たちは全員が女子に招待状を送ります。ただし、各女子が受け入れられる招待状は一通だけです。このとき、女子が実際に受け入れられる招待状の総数を求めます。

入力はm × nの行列として与えられ、セル(i, j)は「男子iが女子jに招待状を送ったかどうか」を表します。値が1なら招待状を送ったこと、0なら送っていないことを意味します。

入力例

この場合の出力は3になります。

• 女子1が男子1の招待状を受け入れる
• 女子2が男子3の招待状を受け入れる
• 女子3が男子2の招待状を受け入れる

(ここではインデックスは1から始まるものとします)

解き方のアプローチ

この問題は「二部グラフの最大マッチング」問題として捉えることができます。男子と女子をそれぞれ頂点とみなし、招待状を辺とみなすと、受け入れられる招待状の最大数は最大マッチングのサイズに一致します。

これを解くには、DFS(深さ優先探索)を使って増加道(オーギュメンティングパス)を見つける手法が有効です。手順は以下の通りです。

1. dfs(node, seen) 関数を定義する。
   • nei を 0 から N-1 まで繰り返す:
     • grid[node][nei] が 0 以外で、かつ seen[nei] が False の場合:
       • seen[nei] := True
       • matching[nei] が -1 である、または dfs(matching[nei], seen) が True の場合:
         • matching[nei] := node
         • True を返す
   • False を返す
2. M := grid の行数
3. N := grid の列数
4. matching := サイズ N のリスト(すべて -1 で初期化)
5. res := 0
6. i を 0 から M-1 まで繰り返す:
   • seen := サイズ N のリスト(すべて False で初期化)
   • dfs(i, seen) が True の場合: res += 1
7. res を返す

実装例

以下のPythonコードで実際の動作を確認できます。

def solve(grid):
    M, N = len(grid), len(grid[0])
    matching = [-1] * N

    def dfs(node, seen):
        for nei in range(N):
            if grid[node][nei] and not seen[nei]:
                seen[nei] = True
                if matching[nei] == -1 or dfs(matching[nei], seen):
                    matching[nei] = node
                    return True
        return False

    res = 0
    for i in range(M):
        seen = [False] * N
        if dfs(i, seen):
            res += 1

    return res

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

入力

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

出力

3

計算量について

各男子ごとにDFSを1回実行する構成のため、このアルゴリズムの時間計算量は一般的にO(V・E)(Vは頂点数、Eは辺数)と評価されます。人数がそれほど多くない場合は十分に高速に動作し、競技プログラミングなどでも広く使われる定番の手法です。