AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC165 C - Many Requirements

問題原文

atcoder.jp

解法

 N M の制約が小さいことに注目する。

全パターン試すのが楽そうだけど、条件を満たすような数列はどれぐらいあるかを考える。
例えばバラバラの数字  k 個を適当に並べてそれが単調増加になっている確率は  \frac{1}{k!}
この事実を考えると最大ケースでもめちゃくちゃ通り数が削られることがわかる。

 Q の制約も小さいので、これら全てについて愚直にスコアを計算しても余裕で間に合う。

けんちょんさんの考え方に勉強させていただきました。ありがとうございます。 drken1215.hatenablog.com

実装

再帰でかける。
.append() -> dfs() -> .pop() とけんちょんさんの記事ではやっていて、なるほどと思った。
(↓ はそれではないです)

from itertools import combinations_with_replacement
N, M, Q = map(int, input().split())
X = [list(map(int, input().split())) for _ in range(Q)]
ans = 0

def calc(arr):
    score = 0
    for a, b, c, d in X:
        if arr[b - 1] - arr[a - 1] == c:
            score += d
    return score


def dfs(arr):
    global ans
    if len(arr) == N:
        ans = max(ans, calc(arr))

    elif len(arr) == 0:
        for nx in range(1, M + 1):
            dfs(arr + [nx])

    else:
        for nx in range(arr[-1], M + 1):
            dfs(arr + [nx])


dfs([])
print(ans)

ところでpythonだと itertools.combinations_with_replacement がめちゃくちゃ便利。
Twitterで教えてもらった(ありがとうございます)

itertools --- 効率的なループ実行のためのイテレータ生成関数 — Python 3.8.3 ドキュメント

中でソートしてる分ちょっと遅いはずだけど、まぁ全然大丈夫。

from itertools import combinations_with_replacement
N, M, Q = map(int, input().split())
X = [list(map(int, input().split())) for _ in range(Q)]

ans = 0
for pattern in combinations_with_replacement(range(1, M + 1), N):
    tmp_ans = 0
    pattern = list(pattern)
    for a, b, c, d in X:
        if pattern[b - 1] - pattern[a - 1] == c:
            tmp_ans += d
    ans = max(ans, tmp_ans)

print(ans)

感想

たくさん知見を得られた良い問題でした。
正答率が思ったより低くてへぇ〜ってなった(数学系の問題の方が強い人多いのかな?)