結果

問題 No.391 CODING WAR
コンテスト
ユーザー 草苺奶昔
提出日時 2023-04-02 20:49:00
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 149 ms / 2,000 ms
コード長 6,486 bytes
コンパイル時間 206 ms
コンパイル使用メモリ 82,312 KB
実行使用メモリ 72,624 KB
最終ジャッジ日時 2024-09-25 00:38:09
合計ジャッジ時間 3,002 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 4
other AC * 16
権限があれば一括ダウンロードができます

ソースコード

diff # 

class 写像十二相:
    """https://qiita.com/drken/items/f2ea4b58b0d21621bd51"""

    __slots__ = ("_fac", "_ifac", "_inv", "_mod")

    def __init__(self, size: int, mod: int) -> None:
        self._mod = mod
        self._fac = [1]
        self._ifac = [1]
        self._inv = [1]
        self._expand(size)

    def query(
        self,
        n: int,
        k: int,
        *,
        isBallSame: bool,
        isBoxSame: bool,
        atMostOneBallPerBox=False,
        noLimitWithBox=False,
        atLeastOneBallPerBox=False,
    ) -> int:
        """n个球放入k个盒子的方案数.

        Args:
            isBallSame (bool): 球是否有区别.
            isBoxSame (bool): 盒子是否有区别.
            atMostOneBalPerBox (bool, optional): 每个盒子最多放一个球.
            noLimitWithBox (bool, optional): 每个盒子可以放任意个球.
            atLeastOneBallPerBox (bool, optional): 每个盒子至少放一个球.
        """
        limits = (atMostOneBallPerBox, noLimitWithBox, atLeastOneBallPerBox)
        assert limits.count(True) == 1, "Must have one limit and only one limit with box."
        if isBallSame and isBoxSame:
            if atMostOneBallPerBox:
                return self._solve1(n, k)
            if noLimitWithBox:
                return self._solve2(n, k)
            if atLeastOneBallPerBox:
                return self._solve3(n, k)
        if not isBallSame and isBoxSame:
            if atMostOneBallPerBox:
                return self._solve4(n, k)
            if noLimitWithBox:
                return self._solve5(n, k)
            if atLeastOneBallPerBox:
                return self._solve6(n, k)
        if isBallSame and not isBoxSame:
            if atMostOneBallPerBox:
                return self._solve7(n, k)
            if noLimitWithBox:
                return self._solve8(n, k)
            if atLeastOneBallPerBox:
                return self._solve9(n, k)
        if not isBallSame and not isBoxSame:
            if atMostOneBallPerBox:
                return self._solve10(n, k)
            if noLimitWithBox:
                return self._solve11(n, k)
            if atLeastOneBallPerBox:
                return self._solve12(n, k)

        raise Exception("Unreachable code.")

    def _solve1(self, n: int, k: int) -> int:
        """有区别的球放入有区别的盒子(每个盒子最多放一个球)."""
        return self.P(n, k)

    def _solve2(self, n: int, k: int) -> int:
        """有区别的球放入有区别的盒子(每个盒子可以放任意个球)."""
        return pow(k, n, self._mod)

    def _solve3(self, n: int, k: int) -> int:
        """有区别的球放入有区别的盒子(每个盒子至少放一个球).
        容斥原理:用总方案数减去不合法的方案数.
        O(k*logn)
        """
        mod = self._mod
        res = 0
        for i in range(k + 1):
            if (k - i) & 1:
                res -= self.C(k, i) * pow(i, n, mod)
            else:
                res += self.C(k, i) * pow(i, n, mod)
            res %= mod
        return res

    def _solve4(self, n: int, k: int) -> int:
        """无区别的球放入有区别的盒子(每个盒子最多放一个球)."""
        return self.C(n, k)

    def _solve5(self, n: int, k: int) -> int:
        """无区别的球放入有区别的盒子(每个盒子可以放任意个球)."""
        return self.C(n + k - 1, n)

    def _solve6(self, n: int, k: int) -> int:
        """无区别的球放入有区别的盒子(每个盒子至少放一个球)."""
        return self.C(n - 1, k - 1)

    def _solve7(self, n: int, k: int) -> int:
        """有区别的球放入无区别的盒子(每个盒子最多放一个球)."""
        return 0 if n > k else 1

    def _solve8(self, n: int, k: int) -> int:
        """有区别的球放入无区别的盒子(每个盒子可以放任意个球).
        贝尔数.
        """
        ...

    def _solve9(self, n: int, k: int) -> int:
        """有区别的球放入无区别的盒子(每个盒子至少放一个球).
        第二类斯特林数.
        """
        ...

    def _solve10(self, n: int, k: int) -> int:
        """无区别的球放入无区别的盒子(每个盒子最多放一个球)."""
        ...

    def _solve11(self, n: int, k: int) -> int:
        """无区别的球放入无区别的盒子(每个盒子可以放任意个球).
        分割数.
        """
        ...

    def _solve12(self, n: int, k: int) -> int:
        """无区别的球放入无区别的盒子(每个盒子至少放一个球).
        分割数.
        """
        ...

    def fac(self, k: int) -> int:
        self._expand(k)
        return self._fac[k]

    def ifac(self, k: int) -> int:
        self._expand(k)
        return self._ifac[k]

    def inv(self, k: int) -> int:
        self._expand(k)
        return self._inv[k]

    def C(self, n: int, k: int) -> int:
        if n < 0 or k < 0 or n < k:
            return 0
        mod = self._mod
        return self.fac(n) * self.ifac(k) % mod * self.ifac(n - k) % mod

    def P(self, n: int, k: int) -> int:
        if n < 0 or k < 0 or n < k:
            return 0
        mod = self._mod
        return self.fac(n) * self.ifac(n - k) % mod

    def H(self, n: int, k: int) -> int:
        """可重复选取元素的组合数"""
        return self.C(n + k - 1, k)

    def put(self, n: int, k: int) -> int:
        """n个相同的球放入k个不同的盒子(盒子可放任意个球)的方案数."""
        return self.C(n + k - 1, n)

    def _expand(self, size: int) -> None:
        if len(self._fac) < size + 1:
            mod = self._mod
            preSize = len(self._fac)
            diff = size + 1 - preSize
            self._fac += [1] * diff
            self._ifac += [1] * diff
            self._inv += [1] * diff
            for i in range(preSize, size + 1):
                self._fac[i] = self._fac[i - 1] * i % mod
            self._ifac[size] = pow(self._fac[size], mod - 2, mod)  # !modInv
            for i in range(size - 1, preSize - 1, -1):
                self._ifac[i] = self._ifac[i + 1] * (i + 1) % mod
            for i in range(preSize, size + 1):
                self._inv[i] = self._ifac[i] * self._fac[i - 1] % mod


MOD = int(1e9 + 7)
X = 写像十二相(int(1e5), MOD)
if __name__ == "__main__":
    n, k = map(int, input().split())
    print(X.query(n, k, isBallSame=True, isBoxSame=True, atLeastOneBallPerBox=True))
0