コード
"""Modint
・つかい方
mint = Modint(998244353)
value = mint(10)
"""
class Modint():
def __init__(self, mod: int):
self.mod: int = mod
def __call__(self, num: int):
return NumberMint(num, self.mod)
class NumberMint():
def __init__(self, num: int, mod: int):
self.num: int = num
self.mod: int = mod
if self.num < 0 or self.mod <= self.num:
self.num %= self.mod
self.__inverse = None # 逆数
def __int__(self) -> int:
return self.num
def __str__(self) -> str:
return str(self.num)
def __repr__(self) -> str:
return "Modint({})".format(self.num)
def __add__(self, other):
if isinstance(other, NumberMint):
new_value = self.num + other.num
else:
new_value = self.num + other
if self.mod <= new_value:
new_value -= self.mod
return NumberMint(new_value, self.mod)
def __sub__(self, other):
if isinstance(other, NumberMint):
new_value = self.num - other.num
else:
new_value = self.num - other
if new_value < 0:
new_value += self.mod
return NumberMint(new_value, self.mod)
def __mul__(self, other):
if isinstance(other, NumberMint):
new_value = self.num * other.num % self.mod
else:
new_value = other % self.mod * self.num % self.mod
return NumberMint(new_value, self.mod)
def __truediv__(self, other):
if isinstance(other, NumberMint):
new_value = self.num * other.inverse % self.mod
else:
other_inverse = pow(other, self.mod - 2, self.mod)
new_value = self.num * other_inverse % self.mod
return NumberMint(new_value, self.mod)
def __floordiv__(self, other):
return self / other
def __pow__(self, power):
return pow(self.num, power, self.mod)
def __iadd__(self, other):
return self + other
def __isub__(self, other):
return self - other
def __imul__(self, other):
return self * other
def __itruediv__(self, other):
return self / other
def __ifloordiv__(self, other):
return self / other
def __radd__(self, other):
return self + other
def __rsub__(self, other):
return NumberMint(other, self.mod) - self
def __rmul__(self, other):
return self * other
def __rtruediv__(self, other):
return NumberMint(other, self.mod) / self
def __rfloordiv__(self, other):
return NumberMint(other, self.mod) / self
def __eq__(self, other) -> bool:
if isinstance(other, NumberMint):
return self.num == other.num
return self.num == other % self.mod
def __ne__(self, other) -> bool:
if isinstance(other, NumberMint):
return self.num != other.num
return self.num != other % self.mod
def __inverse_generator(self):
return pow(self.num, self.mod - 2, self.mod)
@property
def inverse(self) -> int:
if self.__inverse is None:
self.__inverse = self.__inverse_generator()
return self.__inverse
つかい方
インスタンス作成
mint = Modint(998244353) # MODをセットする。
x = mint(15) # これ以降は、xを使った演算はすべて998244353を法にしてくれる。
できる演算の一覧
mint = Modint(998244353)
x = mint(10)
y = mint(7)
print(x + y) # 17
print(x - y) # 3
print(x * y) # 70
print(x / y) # 570425346
print(x // y) # 570425346
print(y ** 5) # 16807
print(x == y) # False
print(x != y) # True
# こういうタイプのもできます！
x += y
x -= y
x *= y
x /= y
x //= y
# 片方だけでもModintだったら、演算結果もModintになります！
print(x + 5) # 15
print(22 + y) # 29
# 逆数
print(x.inverse) # 99824436
コメント
わざわざ % mod を書く手間がはぶけます！
割り算(分数)が直感的に書けます！
mod $998244353$ や mod $10^9+7$ が問題で出てきたら、とりあえずこれを使っておけばmodのことを気にせずに問題に集中できます。
気に入ったら、ぜひつかってみてね。
こだわりポイント (オタク向け)
逆数に関しては、実はクラスの中ではself.__inverse という変数にして外部アクセスができないようにしていて、初期値はNoneにしてあります。
inverseメソッドを使うと、もしself.__inverse がNoneのときは下のようにself.__inverseをちゃんと計算します。Noneでなければ、その値をそのまま返します。
self.__inverse = pow(self.num, self.mod - 2, self.mod)
.inverse() と書くとメソッドの雰囲気が出ちゃうので、propertyデコレータを使うことで変数を扱っている感じにしました。
@property
def inverse(self) -> int:
if self.__inverse is None:
self.__inverse = self.__inverse_generator()
return self.__inverse
問題によってはたくさん逆数が呼び出されることがあって、その度に逆数を計算してたらTLEしちゃうこともあるので、逆数をクラス内の変数に持っておく工夫をしました。