Pythonで「控えめな行列」の配置パターン数を数えるプログラムの作成方法

問題概要

2つの整数 n と m が与えられたとき、n × m の「控えめな行列(humble matrix)」として成立しうる配置の総数を求めます。ある行列が控えめな行列とみなされるのは、以下の条件を満たす場合です。

• 1 から n × m までの各要素が、それぞれちょうど1回ずつ含まれていること
• 任意の2つのインデックスの組 (i1, j1) と (i2, j2) について、(i1 + j1) < (i2 + j2) であるならば、Mat[i1, j1] < Mat[i2, j2] が必ず成り立つこと

答えは非常に大きな値になる可能性があるため、結果は 10^9 + 7 で割った余りとして返します。

入出力例

たとえば n = 2、m = 2 の場合、出力は 2 になります。条件を満たす行列は次の2通りしか存在しないためです。

および

解法の考え方

この問題の鍵となるのは、「i + j の値が等しいセル(反対角線上のセル)同士には大小関係の制約がない」という点です。一方で、i + j が小さいセルの値は、i + j が大きいセルのすべての値よりも小さくなければなりません。その結果、各反対角線に割り当てられる値は、全要素を昇順に並べたときの連続した区間へと固定されます。

つまり、サイズが s である反対角線には内部の並べ方として s! 通りの自由度があり、求める答えは「各反対角線のサイズの階乗の積」で表されます。n ≥ m とおくと、反対角線のサイズは 1, 2, …, m−1, m, …, m, m−1, …, 2, 1 という形になり(中央の m は n − m + 1 回繰り返される)、答えは次の式に帰着します。

(1! × 2! × … × (m−1)!)² × m!^n−m+1

この式を高速に計算するために、あらかじめ 10^6 までの階乗を mod 10^9 + 7 で前計算しておくのがポイントです。

アルゴリズムの手順

• p := 10^9 + 7(剰余を取るための素数)
• result := 値 1 を持つリスト(階乗テーブル)
• x を 2 から 10^6 までループし、「result の末尾の要素 × x mod p」を result の末尾に追加していく
• m > n であれば、n と m の値を入れ替える
• prod := 1 とする
• x を 1 から m − 1 までループし、prod に result[x−1](すなわち x!)を掛けて mod p を取る
• prod を2乗する
• x を 0 から n − m までループし、prod に result[m−1](すなわち m!)を掛けて mod p を取る
• prod を返す

Pythonでの実装例

それでは、実際の実装を見ながら理解を深めましょう。

p = 10**9+7

def solve(n, m):
   result = [1]
   for x in range(2,10**6+1):
      temp = result[-1]
      temp = (temp*x) % p
      result.append(temp)

   if(m > n):
      temp = n
      n = m
      m = temp
   prod = 1
   for x in range(1,m):
      prod = (prod * result[x-1]) % p
   prod = (prod**2) % p
   for x in range(n-m+1):
      prod = (prod*result[m-1]) % p
   return prod

n = 3
m = 3
print(solve(n, m))

入力

3, 3

出力

24

n = m = 3 の場合、反対角線のサイズは 1, 2, 3, 2, 1 となるため、答えは 1! × 2! × 3! × 2! × 1! = 24 となり、プログラムの出力と一致します。