-
Notifications
You must be signed in to change notification settings - Fork 129
/
combinations.py
68 lines (62 loc) · 2.32 KB
/
combinations.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
class Factorial:
def __init__(self, MOD):
self.MOD = MOD
self.factorials = [1, 1]
self.invModulos = [0, 1]
self.invFactorial_ = [1, 1]
def calc(self, n):
if n <= -1:
print("Invalid argument to calculate n!")
print("n must be non-negative value. But the argument was " + str(n))
exit()
if n < len(self.factorials):
return self.factorials[n]
nextArr = [0] * (n + 1 - len(self.factorials))
initialI = len(self.factorials)
prev = self.factorials[-1]
m = self.MOD
for i in range(initialI, n + 1):
prev = nextArr[i - initialI] = prev * i % m
self.factorials += nextArr
return self.factorials[n]
def inv(self, n):
if n <= -1:
print("Invalid argument to calculate n^(-1)")
print("n must be non-negative value. But the argument was " + str(n))
exit()
p = self.MOD
pi = n % p
if pi < len(self.invModulos):
return self.invModulos[pi]
nextArr = [0] * (n + 1 - len(self.invModulos))
initialI = len(self.invModulos)
for i in range(initialI, min(p, n + 1)):
next = -self.invModulos[p % i] * (p // i) % p
self.invModulos.append(next)
return self.invModulos[pi]
def invFactorial(self, n):
if n <= -1:
print("Invalid argument to calculate (n^(-1))!")
print("n must be non-negative value. But the argument was " + str(n))
exit()
if n < len(self.invFactorial_):
return self.invFactorial_[n]
self.inv(n)
nextArr = [0] * (n + 1 - len(self.invFactorial_))
initialI = len(self.invFactorial_)
prev = self.invFactorial_[-1]
p = self.MOD
for i in range(initialI, n + 1):
prev = nextArr[i - initialI] = (prev * self.invModulos[i % p]) % p
self.invFactorial_ += nextArr
return self.invFactorial_[n]
class Combination:
def __init__(self, MOD):
self.MOD = MOD
self.factorial = Factorial(MOD)
def ncr(self, n, k):
if k < 0 or n < k:
return 0
k = min(k, n - k)
f = self.factorial
return f.calc(n) * f.invFactorial(max(n - k, k)) * f.invFactorial(min(k, n - k)) % self.MOD