ABC473Dを解いた【DFS全探索】

筆者はレート800前後の茶～緑コーダ

ABC473のD問題を解いていく

実装コード

1 から N までの各個数を DFS で決め、重み付き和が K になる列をすべて列挙する。

• N, K を受け取り、各値の個数を保存する配列 path を用意する
• dfs(i, cur) で、1 から i までに決めた個数による重み付き和を cur として管理する
• 値 i + 1 の個数 j を、合計が K を超えない範囲で試す
• path[i] に j を入れ、cur + j * (i + 1) を渡して次の値へ進む
• 最後の値 N まで来たら、残りの K - cur が N で割り切れるか確認する
• 割り切れる場合は最後の個数を path[-1] に入れ、条件を満たす path を出力する

main.py

from bisect import bisect_left, bisect_right, insort_left, insort_right
from collections import defaultdict, Counter, deque
from functools import reduce, lru_cache
from itertools import product, accumulate, groupby, combinations
import sys
import os
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rI1(): return (int(sys.stdin.readline().rstrip())-1)
def rLI1(): return list(map(lambda a:int(a)-1,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
IS_LOCAL = int(os.getenv("ATCODER", "0"))==0
err = (lambda *args, **kwargs: print(*args, **kwargs, file=sys.stderr)) if IS_LOCAL else (lambda *args, **kwargs: None)

def main():
    N, K = rLI()
    path = [0] * N

    def dfs(i, cur):
        if i + 1 == N:
            if (K - cur) % N == 0:
                path[-1] = (K - cur) // N
                print(*path)
            return

        for j in range((K - cur) // (i + 1) + 1):
            path[i] = j
            dfs(i + 1, cur + j * (i + 1))

    dfs(0, 0)

if __name__ == '__main__':
    main()

感想

最後の要素をループで試さず、残りが N で割り切れるときだけ個数を確定することで、全探索を簡潔に書けるのが参考になった。
相変わらずDFSを再帰で実装するのはなかなか慣れない。