[AtCoder]ABC165C個人的メモ[Python]

ABC165 C - Many Requirements

考えられる数列Aをdfsによって全探索し，得点の最大値を求める．
解答例を2つ示す

• 解答例1-dfsを再帰関数で実装

  • calc関数
    数列Aから得点の計算をする関数

  • dfs関数
    考えられる数列Aをdfsによって全探索しつつ，得点の最大値を求める関数
    数列Aの長さがNなら得点を計算した後，今までの最高得点と比較，最高得点を記録
    それ以下の長さなら，数列Aの制約を満たす様に要素を1つ追加して再帰
    Aが空リストなら$x=1$を，それ以外なら$A[-1]\leqq x\leqq M$の$x$をAに追加して再帰

  main.py

    #!/usr/bin/env python3
    def main():
        import sys
  
        def calc(A: list):
            score = 0
            for a, b, c, d in lst:
                if A[b - 1] - A[a - 1] == c:
                    score += d
            return score
  
        def dfs(A: list):
            nonlocal ans
            if len(A) == N:
                ans = max(ans, calc(A))
                return
            for v in range(A[-1] if A else 1, M + 1):
                A.append(v)
                dfs(A)
                A.pop()
  
        input = sys.stdin.readline
  
        N, M, Q = map(int, input().split())
        lst = [list(map(int, input().split())) for _ in range(Q)]
  
        ans = 0
        dfs([])
        print(ans)
  
  
    main()
  
  

• 解答例2-itertools.combinations_with_replacement(iterable, r)
  itertools.combinations_with_replacementをつかって，数列Aを全探索する
  公式ドキュメントによると，itertools.combinations_with_replacement(iterable, r)は，入力iterableから，それぞれの要素が複数回現れることを許して長さrの要素の部分列をタプルで返す

  test.py

  # itertools.combinations_with_replacement(iterable, r)の使用例
  from itertools import combinations_with_replacement
  
  for pattern in combinations_with_replacement('ABCD', 2):
      print(pattern, end='')
  
  # AA AB AC AD BB BC BD CC CD DD
  
  

  上記の関数を用いると，解答は以下の様に書ける

  main.py

    #!/usr/bin/env python3
    def main():
        import sys
        from itertools import combinations_with_replacement
  
        input = sys.stdin.readline
  
        N, M, Q = map(int, input().split())
        lst = [list(map(int, input().split())) for _ in range(Q)]
  
        ans = 0
        for pattern in combinations_with_replacement(range(1, M + 1), N):
            res = 0
            pattern = list(pattern)
            for a, b, c, d in lst:
                if pattern[b - 1] - pattern[a - 1] == c:
                    res += d
            ans = max(ans, res)
  
        print(ans)
  
  
    main()
  
  

参考URL

• あっとのTECH LOG
  ABC165 C - Many Requirements
  https://at274.hatenablog.com/entry/2020/05/20/213705
• けんちょんの競プロ精進記録
  よくやる再帰関数の書き方 〜 n 重 for 文を機械的に 〜
  https://drken1215.hatenablog.com/entry/2020/05/04/190252
• Python 3.8.5 ドキュメント
  効率的なループ実行のためのイテレータ生成関数
  https://docs.python.org/ja/3/library/itertools.html#itertools.combinations_with_replacement