結果

問題 No.117 組み合わせの数
ユーザー Daisuke MiyakawaDaisuke Miyakawa
提出日時 2022-06-18 10:22:44
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 740 ms / 5,000 ms
コード長 2,559 bytes
コンパイル時間 362 ms
コンパイル使用メモリ 82,280 KB
実行使用メモリ 266,684 KB
最終ジャッジ日時 2024-10-09 21:40:21
合計ジャッジ時間 2,245 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 1
権限があれば一括ダウンロードができます

ソースコード

diff # 

class ModComb:
    """\
    N! mod p 、またそれを前提とした nCr, nPr, nHr の計算を前処理付きで高速に計算する
    nHr = (n+r-1)Cr であるため、重複組合せを求める際にはNの数値を2倍程度にする必要がある可能性に注意
    """
    def __init__(self, N, p: int = 1_000_000_007):
        """\
        :param N: 乗数の最大値。
        :param p: 剰余に用いる数。このクラスが正しく機能するためには素数である必要がある
        """
        self._N = N
        self._p = p
        self._factorial_cache = {0: 1}
        for i in range(1, N + 1):
            self._factorial_cache[i] = (self._factorial_cache[i - 1] * i) % p

    def factorial(self, n):
        return self._factorial_cache[n]

    def inverse(self, x):
        """xの逆数 (mod p) を返す"""
        # フェルマーの小定理からxの逆数をx'とすると x' = x^(p-2) mod p
        k = self._p - 2
        ret = 1
        y = x
        while k:
            if k & 1:
                ret = (ret * y) % self._p
            y = (y * y) % self._p
            k //= 2
        return ret

    def P(self, n, r) -> int:
        """nPr mod p を求める"""
        if n < r:
            return 0
        a = self.factorial(n)
        b = self.factorial(n - r)
        return (a * self.inverse(b)) % self._p

    def C(self, n, r) -> int:
        """nCr mod p を求める"""
        if n < r:
            return 0
        a = self.factorial(n)
        b = self.factorial(n - r)
        c = self.factorial(r)
        bc = (b * c) % self._p
        return a * self.inverse(bc) % self._p

    def H(self, n, r) -> int:
        """\
        nHr mod p (重複組合せ) を求める
        """
        if r == 0:
            return 1
        return self.C(n + r - 1, r)


def main():
    import re
    import sys
    if len(sys.argv) > 1:
        args = sys.argv[1:]
    else:
        T = int(input())
        args = [input() for _ in range(T)]
    mc = ModComb(2 * 10 ** 6)
    for arg in args:
        m = re.match(r"([CPH])\((\d+)\s*,\s*(\d+)\)", arg)
        if m:
            n, r = int(m.group(2)), int(m.group(3))
            if m.group(1) == "P":
                print(mc.P(n, r))
                continue
            elif m.group(1) == "C":
                print(mc.C(n, r))
                continue
            elif m.group(1) == "H":
                print(mc.H(n, r))
                continue
        raise RuntimeError(f"Unknown format \"{arg}\"")


if __name__ == "__main__":
    main()
0