めもめも

このブログに記載の内容は個人の見解であり、必ずしも所属組織の立場、戦略、意見を代表するものではありません。

O - Matching の解説(その1)

何の話かと言うと

atcoder.jp

上記の問題をネタに・・・うーん。これ、個人的には、「いやよくこんな方法思いついたな」系で、あまりパターン化できないんですが、とりあえずは、できるだけ天下り的な説明はさけて、論理的に正解に近づいていく方法で説明してみます。

まずは「しらみつぶし」で考える

はい。何をやっていいか皆目検討がつかない時は、まずは、「しらみつぶし」で考え方のヒントを探りましょう。

男性を 1 人目から N 人目までスキャンしながら、それぞれにペアになる女性を選択していくという再帰的なループが想像できます。ある男性と女性のペアを選択したら、今度は残った男性群、女性群について再帰的に同じ問題を解けばよいわけです。そこで、

・N 人の男性から「今残っている男性」を示す変数 males ---- (1)
・N 人の女性から「今残っている女性」を示す変数 females ---- (2)

を用意してこれらを引数とする関数 solve(males, females) を再帰的に呼び出す方法を試してみましょう。

(1)(2) の変数のデータ構造は、N 個の True/False を並べたタプルなどが考えられますが、ここでは、(競技プログラミングっぽく?)2 進数を用いたバイナリー表現を使うことにしましょう。N ビットの整数値、すなわち、0 \sim 2^N-1 の範囲の整数値を用いて、下から n ビット目が 1 であれば、n 番目の人が残っているということにします。

そうすると、まずは、こんな感じのプロトタイプができあがります。

def solve(males, females):  # 与えられた男性群、女性群でできる完全ペアの場合の数を返す
  global T, bitcnt   # T : matching table
  mod = 10**9 + 7
  N = len(T)

  # todo: ボトムケースを実装
 
 count = 0
  for m in range(N):
    if not (males & 1<<m): # m 人目の男性がいるか?
      continue
    for f in range(N):
      if not (females & 1<<f): # f 人目の女性がいるか?
        continue
      if T[m][f] == 1:
        males2 = males ^ 1<<m     # m 人目の男性を削除
        females2 = females ^ 1<<f # f 人目の女性を削除
        count += solve(males2, females2)
        count %= mod
  return count

関数 solve(males, females) は与えられた男性群、女性群でできる完全ペアの場合の数を返します。ここで「完全ペア」と言っているのは、全員マッチできずに必ず誰かが残るというケースが考えられるので、そういう場合は除く(つまり、一人も残らずにマッチできる場合を考える)ことを示します。

ここでは、すべての男性と女性のペアを網羅するループで、実際にマッチできる組みあわせの場合は、この2人を除いた後の場合の数を再帰的に計算して加えていきます。

完全ペアとなるのは、男性も女性も 0 人になる場合ですので、ボトムケースは次の様に埋められます。

  # todo: ボトムケースを実装
  size1 = bitcnt[males]   # 男性の総数
  size2 = bitcnt[females] # 女性の総数
  if size1 == 0 and size2 == 0: # ペア完成に成功
    return 1

bitcnt[m] はバイナリ表現 m に含まれる人数(ビットが 1 の個数)を事前に計算したリストです。また、男性と女性が同数の状態からペアで引いていくので、論理的には、どちらか一方が 0 になっていることだけチェックすれば十分ですが、ここでは、明示的に両方をチェックしています。(人数が 0 であることをチェックするだけであれば、「males == 0 and females == 0」という条件でOKなのですが、後で使うことを考えて、0 以外の場合も bitcnt[ ] で調べられるようにしてあります。)

これで、まずは動くものができましたが、これを実行すると、期待される場合の数よりも大きな値がでてきてしまいます。どこかでダブルカウントをしているんですよね・・・・。

print デバッグをがんばるとわかる様に、男性と女性をマッチする順番から重複が生じています。たとえば、男女それぞれ2人ずつで、

(1, 1), (2, 2)

という1つの場合しか答えがない場合、上記の2重ループだと、

(1, 1), (2, 2)
(2, 2), (1, 1)

というように、ペアをつくる順番を入れ替えたものがダブルカウントされてしまいます。solve() の最初の呼び出し(再帰呼び出しが行われる前)で男性をループしていますが、このループの中で、1人目を先にマッチして再帰する場合と、2人目を先にマッチして再帰する場合の両方がおきているんですよね。

これを解消するには、男性すべてをループする必要はなくて、与えられたグループ male の中の最初の男性だけをマッチして、2人目以降は再帰処理にまかせればよいはずです。

def solve(males, females):
  global T, bitcnt   # T : matching table
  mod = 10**9 + 7
  N = len(T)
 
  size1 = bitcnt[males]   # 男性の総数
  size2 = bitcnt[females] # 女性の総数
  if size1 == 0 and size2 == 0: # ペア完成に成功
    return 1

  male_match = False
  count = 0
  m = 0
  while True:  # グループ内の最初の男性を検索
    if males & 1<<m:
      break
    m += 1
  for f in range(N):
    if not (females & 1<<f): # f 人目の女性がいるか?
      continue
    if T[m][f] == 1:
      males2 = males ^ 1<<m     # m 人目の男性を削除
      females2 = females ^ 1<<f # f 人目の女性を削除
      count += solve(males2, females2)
      count %= mod
  return count

これで、本当に正しく動くものができました。再帰関数をメモ化して高速化したものを提出してみましょう。

atcoder.jp

結果は、約半数のケースで TLE(実行時間オーバー)です。素朴な「しらみつぶし」としては悪くない結果かも知れません。

「しらみつぶし」からアイデアを得る

さあ、さきほどの「しらみつぶし」案から、DPっぽい実装のヒントは得られたでしょうか・・・? 男性をマッチする際は、1人目から順番にマッチしていけばOKでした。ということは・・・

・dp[1] : 1人目の男性だけが参加する場合の答え
・dp[2] : 1人目と2人目の男性が参加する場合の答え
 \vdots
・dp[N] : 1人目〜N人目の男性(要するに全員)が参加する場合の答え

というように、「何人目までの男性を参加させるか」を問題のサイズとして、サイズの小さい方から順に解いていくことができるやも知れません。

先ほどの再帰関数の動き(再帰呼び出しの流れ)を想像すると、初回の実行(再帰呼び出しが行われる前)で1 人目の男性(男性1)にマッチさせる女性を選んで、選ばれたそれぞれの女性について、(再帰呼び出しによって)その女性を除いた中から、2人目の男性(男性2)にマッチする女性を選ぶ、という流れになっています。つまり、2人目の男性についての計算(つまり dp[2] の計算)をする際は、「男性1がどの女性を選ばなかった」という追加情報が必要になります。

うーん。。。。わかりにくいですね。具体例でいきましょう。

・dp[2] : 1人目と2人目の男性が参加する前提での場合の数

を計算しているとして、2人目の男性が女性 f を選択した場合を考えると、1人目の男性はこの女性 f を除く女性群に対してマッチしているはずなので・・・

・dp[1][females]:女性群 females に対して、男性1がマッチする場合の数(男性がすべてマッチできればよくて、女性が余るのはOKとする)

があらゆる組み合わせの females(つまり、females = 1 \sim 2^N-1)に対して計算できていれば、

・dp[2][females]:女性群 females に対して、男性1と男性2がマッチする場合の数(男性がすべてマッチできればよくて、女性が余るのはOKとする)

については、

for f in # [females の中で、男性2とマッチ可能な女性]
  dp[2][females] += dp[1][females ^ 1<<f] # (females ^ 1<<f) は、females から女性 f を除いたグループ

というループで計算できるはずです。

男性を増やしながらこの計算を進めれば、最後は、すべての男性が参加した場合の数として答えが得られそうです。やってみましょう。

import sys

def main(f):
  mod = 10**9 + 7
  N = int(f.readline())
  T = [None]
  for i in range(N):
    T.append([None] + list(map(int, f.readline().split())))

  dp = [[0] * 2**N for _ in range(N+1)] # dp[n][females] : 女性群 females に対して男性 1 〜男性 n をマッチする場合の数

  # n = 1 の場合
  for females in range(2**N): # 女性についてはあらゆる組み合わせを網羅する
    for f in range(1, N+1): # 男性 1 とマッチする女性を探す
      if females & 1<<(f-1) and T[1][f] == 1: # 女性 f がグループにいてマッチ可能
        dp[1][females] = 1

  # n = 2, 3, ..., N の場合
  for n in range(2, N+1):
    for females in range(2**N): # 女性についてはあらゆる組み合わせを網羅する
      for f in range(1, N+1): # 男性 n とマッチする女性を探す
        if females & 1<<(f-1) and T[n][f] == 1: # 女性 f がグループにいてマッチ可能
          dp[n][females] += dp[n-1][females ^ 1<<(f-1)] # 女性 f を除いて、残りを n 未満の男性とマッチする場合の数
          dp[n][females] %= mod

  print(dp[N][2**N])

main(sys.stdin)

これは実際に提出可能なコードです。問題の答えは、dp[N][2**N-1] (男性Nまでをすべての女性 females = 2**N-1 とマッチする場合の数)として得られます。

ただ、残念ながら、多くのケースで TLE になります。方向性としては間違っていませんが、まだちょっと冗長な計算をしているようです。

ここから先のチューニングはまた次回に解説しましょう。(パチパチパチ)