[Python] ABC018 D 100点

問題

D – バレンタインデー

回答

女子の組み合わせ、男子の組み合わせを全て列挙すると\( {}_{18} C _9 \times {}_{18} C _9 \)、女子の組み合わせだけを列挙しその組み合わせに対してチョコをもらう可能性のある男子を考えると \( {}_{18} C _9 \times 18 \) と、計算量を減らすことができます。

AtCoder Beginner Contest 018 解説 from AtCoder Inc.

import itertools

N, M, P, Q, R = map(int, input().split())

happiness = [[0] * (M+1) for _ in range(N+1)]
for _ in range(R):
    x, y, z = map(int, input().split())
    happiness[x][y] = z 

girls_cmb = list(itertools.combinations(range(1, N+1), P))

ans = 0
# グループ内の女子を固定
for girls in girls_cmb:
    # ある男子がグループに入るといくら幸福度が得られるか
    happiness_selected_boy = [0] * (M + 1)
    for girl in girls:
        for boy in range(1, M+1):
            happiness_selected_boy[boy] += happiness[girl][boy]
    # 幸福度の大きい男子を貪欲にQ人選ぶ
    ans = max(ans, sum(sorted(happiness_selected_boy, reverse=True)[:Q]))

print(ans)