comb.py

# Shunta の自作ライブラリ
# https://github.com/NAVYSHUNTA/atcoder-shunta-library/blob/main/math/comb/python/comb.py

# 組合せクラス
class Comb:
    __fac: list[int]
    __fac_inv: list[int]
    __inv: list[int]
    __mod: int | None

    # O(n): コンストラクタ
    def __init__(self, n: int, mod: int | None = None) -> None:
        self.__mod = mod
        self.__fac = [1] * (n + 1)
        self.__fac_inv = [1] * (n + 1)
        self.__inv = [1] * (n + 1)

        if mod is None:
            for i in range(2, n + 1):
                self.__fac[i] = self.__fac[i - 1] * i
        else:
            for i in range(2, n + 1):
                self.__fac[i] = (self.__fac[i - 1] * i) % mod
                self.__inv[i] = (-self.__inv[mod % i] * (mod // i)) % mod
                self.__fac_inv[i] = (self.__fac_inv[i - 1] * self.__inv[i]) % mod

    # nCr の値を求めるメソッド
    # O(1): コンストラクタで mod を指定していないかつ n が小さい場合
    # O(1): コンストラクタで mod を指定している場合（n の値によらない）
    def get_comb(self, n: int, r: int) -> int:
        if n < r or min(n, r) < 0:
            return 0

        if self.__mod is None:
            return self.__fac[n] // (self.__fac[r] * self.__fac[n - r])
        else:
            return (self.__fac[n] * self.__fac_inv[r] * self.__fac_inv[n - r]) % self.__mod