競プロ初中級者100問

ABC021D 多重ループ

この記事で使うアルゴリズム

逆元

はじめに

カテゴリー競プロ初中級者100問では、Qiitaにて@e869120さんがレッドコーダーが教える、競プロ・AtCoder上達のガイドライン【中級編：目指せ水色コーダー！】としてまとめられている100問をPythonで解いています。

全問題の一覧はこちらです

問題

https://atcoder.jp/contests/abc021/tasks/abc021_d

この問題では逆元の計算が求められます。

方針

準備運動として、下記の組み合わせの数え上げについてを考えます。

1
1 <= a_1 < a_2 < ... < a_k <= N

少し考察すると、これはN個の要素の中から重複せずにk個の要素を取り出すパターン数、すなわちn_C_kと等価であることがわかります。

そのように理解すると、本題の

1
1 <= a_1 <= a_2 <= ... <= a_k <= N

はN個の要素の中から重複を許してk個の要素を取り出すパターン数であることがわかります。これは重複組合せと呼ばれる概念で、

1
n_H_k = (n-1+k)_C_k = (n-1+k)!/((n-1)!k!)

を計算すれば良いです。

階乗の部分が巨大になるため、計算途中で剰余の計算を行います。しかしながら、除算が含まれるため逆元と呼ばれる概念を導入する必要があります。

逆元については以下の問題の解説をご覧ください。

ABC034C 散歩

コード

1
2
3
4
5
6
7
8
9
10
11
MOD = 10**9+7
n = int(input())
k = int(input())
MX = 2 * 10**5 # n+k-1が最大となる場合に対応
fact = [1] * (MX+1) # 階乗を格納するリスト
factinv = [1] * (MX+1) # 階乗の逆元を格納するリスト
for i in range(MX):
    fact[i+1] = fact[i] * (i+1) % MOD # 階乗を計算
    factinv[i+1] = pow(fact[i+1], MOD-2, MOD)# MODを法とした逆元
ans = fact[n+k-1] * factinv[k] * factinv[n-1]
print (ans%MOD) # 最後にMODを忘れずに

記事情報

  • 投稿日：2020年5月3日
  • 最終更新日：2020年5月8日