Python3: mod7占い (paizaランク S 相当)やってみた

from itertools import combinations

N = int(input())
DECK = [int(input()) for _ in range(N)]

print(
    sum(1 for cards in combinations(DECK, 3) 
        if sum(cards) % 7 == 0
        )
    )

デッキを読み込んで、3枚足すと7の倍数になる組み合わせを計数していく。まんまですな。

これでもタイムアウトせず、ACしてしまったんだけど、解法例を見ると直接modを取るんじゃなくて、早い段階でmod 7を取って、nCrの計算に持ち込む感じらしい。

そっちの方向でやってみた。

from itertools import combinations_with_replacement, groupby
from math import prod

# n => r => nCr
ncr = (lambda n: lambda r:
       prod(range(n, n - r, -1)) // prod(range(r, 0, -1))
       )

# 0~6のうち、三つ足すと7の倍数になる重複組合せを
# (値, 頻度)のタプルのジェネレーターに
mod7 = ( ((k, len(list(g))) for k, g in groupby(x))
         for x in combinations_with_replacement(range(7), 3)
         if sum(x) % 7 == 0
         )

N = int(input())

a = [0] * 7
for _ in range(N):
    a[int(input()) % 7] += 1

print(
    sum(
        prod( ncr(a[i])(r) for i, r in m)
        for m in mod7
        )
    )

うん、確かに速くなった。