結果

問題 No.117 組み合わせの数
ユーザー はむ吉🐹はむ吉🐹
提出日時 2016-07-25 20:25:21
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 2,862 bytes
コンパイル時間 209 ms
コンパイル使用メモリ 82,432 KB
実行使用メモリ 156,408 KB
最終ジャッジ日時 2024-04-24 09:27:09
合計ジャッジ時間 3,401 ms
ジャッジサーバーID
(参考情報)
judge1 / judge4
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 WA -
権限があれば一括ダウンロードができます

ソースコード

diff #

#!/usr/bin/env pypy3

import array
import re


# Based on http://hos.ac/slides/20130319_enumeration.pdf
class ModPrime(object):

    def __init__(self, p, limit=10**8, typecode="L"):
        assert limit >= 1
        self.p = p
        self.current_limit = 0
        self.typecode = typecode
        self.inv = array.array(typecode, (0, 1))
        self.fact = array.array(typecode, (1,))
        self.fact_inv = array.array(typecode, (1,))
        self.calc(limit)

    def calc(self, new_limit):
        ex_len = new_limit - self.current_limit
        assert ex_len > 0
        self.inv.extend(0 for _ in range(ex_len))
        self.fact.extend(0 for _ in range(ex_len))
        self.fact_inv.extend(0 for _ in range(ex_len))
        for i in range(max(2, self.current_limit), new_limit + 1):
            self.inv[i] = self.p - self.p // i * self.inv[self.p % i] % self.p
        for i in range(max(1, self.current_limit), new_limit + 1):
            self.fact[i] = (self.fact[i - 1] * i) % self.p
            self.fact_inv[i] = (self.fact_inv[i - 1] * self.inv[i]) % self.p
        self.current_limit = new_limit

    def choose(self, n, k):
        if k < 0 or k > n:
            return 0
        elif n < self.p:
            ans = self.fact[n] * self.fact_inv[k]
            ans %= self.p
            ans *= self.fact_inv[n - k]
            ans %= self.p
            return ans
        ans = 1
        while k > 0:
            n0 = n % self.p
            k0 = k % self.p
            if n0 < k0:
                return 0
            ans = ans * self.fact[n0] * self.fact_inv[k0] % self.p
            ans = ans * self.fact_inv[n0 - k0] % self.p % self.p
            n //= self.p
            k //= self.p
            return ans

    def permutation(self, n, k):
        if k < 0 or k > n:
            return 0
        return self.fact[n] * self.fact_inv[n - k]

    def multi_choose(self, n, k):
        if n == 0 and k == 0:
            return 1
        elif n <= 0 or k < 0:
            return 0
        return self.choose(n + k - 1, k)

    def inverse(self, n):
        return self.inv[n]

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

    def factorial_inverse(self, n):
        return self.fact_inv[n]


def main():
    p = 10 ** 9 + 7
    limit = 2 * 10 ** 6 + 10
    mp = ModPrime(p, limit)
    ptn = re.compile(r"(?P<func>[CPH])\((?P<n>[0-9]+),(?P<k>[0-9]+)\)")
    t = int(input())
    for _ in range(t):
        mobj = re.match(ptn, input())
        assert mobj is not None
        f = mobj.group("func")
        n = int(mobj.group("n"))
        k = int(mobj.group("k"))
        if f == "C":
            print(mp.choose(n, k))
        elif f == "P":
            print(mp.permutation(n, k))
        elif f == "H":
            print(mp.multi_choose(n, k))
        else:
            assert False


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