このシリーズではE869120さんによって作成された競プロ典型90問をPythonで解いています。
問題
ポイント
Pythonで組み合わせを生成するにはitertoolsのcombinationsを使用すると便利です。
組み合わせを全探索し、条件に合致するものを見つけていきます。
1点注意したいのが、P で割るとQ 余ること確認する際、組み合わせの整数をa,b,c,d,eとして
(a*b*c*d*e)%P==Qを実行しようとすると、計算の桁が大きくなることで処理に時間がかかり、TLEとなってしまう点です。
ですので、計算途中で桁が大きくなりすぎないように
a % P * b % P * c % P * d % P * e % P == Q で実行します。※上記の式と結果は同じです。
ちなみにPythonではこれでもTLEとなるためPyPy3で提出します。
PythonでもACされる解法が別にあるそうですが、レベル高めなので今回は割愛…。
以上を踏まえ、実装したコードがこちらです。
コード
import sys
from itertools import combinations
sys.setrecursionlimit(10 ** 9)
input = sys.stdin.readline
N, P, Q = map(int, input().split(" "))
nums = list(map(int, input().split(" ")))
ans = 0
combs = combinations(nums, 5)
for a, b, c, d, e in combs:
if a % P * b % P * c % P * d % P * e % P == Q: # (a*b*c*d*e)%PだとTLEになる
ans += 1
print(ans)
