AtCoder Beginners Contest 165のC問題『Many Requirements』をPython3で解く方法を、なるべく丁寧に解説していきます。

とても長くなったので、C問題だけ分割しました。

3種類の方法で解説します。

itertools.combinations_with_replacement()を使う（一番楽）
キュー再帰で作る（汎用性が高い方法）
深さ優先探索（DFS）（解説PDFの方法、めんどくさい）

ABC165C『Many Requirements』

問題ページ：C - Many Requirements
むずかしさ：★★★★★★★★★★（難しいD問題レベル！）
タイプ：総当りの発想、深さ優先探索（他の方法もあり）、過去問演習

まず問題文を読んで意味を理解するのが大変です。意味がわかっても、すべての数列を作って確かめるという発想が必要です。そして、総当りすることがわかっても、どうやって数列を作るか知らないと解けません。

正直言って、難しいD問題の難易度です。

やること

（やばそうなので、D問題が簡単ではないか確認する）
問題文を解読する
解法を考える
実装を考える

ステップ2: 問題文を解読する

問題文をなんとか解読すると、私たちに作ってほしいのは長さが $N$ の数列だそうです。長さ $N$ は2～10です。数列が条件を満たしていると得点をもらえるので、その得点の最大値を出してほしいと言っています。

数列の数字は $1$ 以上 $M$ 以下です。上限 $M$ は1～10です。そして、数字は1個前の数字と同じか大きくないといけません。数字が途中で減ってはいけないということです。

例えば、 $N = 4$ で $M = 3$ だとします。

許される数列の例をあげてみます。
$1,1,2,3$
$1,1,1,1$（すべて同じでもいい）
$2,3,3,3$（$1$からはじまっていなくてもいい）

ダメな数列の例をあげてみます。
$1,1,3,2$（$2$ は $3$ より大きいので、3の後ろには来れません）
$1,2,3,4$（$4$ は上限 $M = 3$ を超えています）
$0,1,2,3$（$1$ 以上なので、$0$ はダメです）
$1,1,1$（長さ $N = 4$です）
$1,1,1,1,1$（同上）

最後に、条件が $Q$ 個あります。条件は最大で50個あって、1つの条件は、4つの整数 $a, b, c, d$ からなります。

その条件とは、

(作った数列の $b$ 番目) - (作った数列の $a$ 番目) = $c$

であることです。もしそうならば、得点が $d$ もらえます。

例えば、条件が
$a=1$
$b=4$
$c=2$
$d=5$
の場合、(数列の $4$ 番目) - (数列の $1$ 番目) = $2$ だと $5$ 点もらえます。

いくつか数列の例をあげると、 $1,1,2,3$ や $1,2,3,3$ は $5$ 点もらえますが、 $1,2,2,2$ や $2,2,3,3$ はもらえません。

これをすべての条件に対して確認して求めた得点の合計の、最大値を求めてほしいそうです。

ステップ3　解法を考える

この問題はあり得る数列を全て作って、それぞれの得点を計算して最大値を求めるしかありません。ですので、総当りできることに気づかないと絶対に解けません。

数列の長さは最大10、数字の上限は最大10なので、作れる数列はそんなに多くない気がします。解説PDFに書いてあるとおりに正確な数を求めると、~~20C10=184756~~ $ _{19} C_{10} = 92378$ 通りになります。（5/3 訂正しました。@forzaMilanさんありがとうございます！）

そして、得点をもらうための条件は最大で50個なので、作った数列に対して全部確認しても間に合いそうな感じがします。

ステップ4 実装を考える

さて、総当りできることに気づいても、どうやって数列をすべて作るかがわからないと解けません。

こういう数列や文字列を作る問題はたまに出てくるので、類題を解いていればわかるかもしれません。しかし、知らなければコンテスト中に思いつくのは難しいと思います。

数列を作る方法はいくつかあります。

itertools.combinations_with_replacementを使う（一番楽）
キュー再帰で作る（汎用性が高い方法）
深さ優先探索（解説PDFの方法）

方法1 combinations_with_replacementを使う方法

一番楽な方法は、itertoolsモジュールの、combinations_with_replacement()関数を使う方法です。私は後から知りましたが、この関数を使うだけで、この問題の数列をすべて列挙することができます。

名前が長くて覚えづらいですが、便利な関数です。使うには、itertoolsモジュールをインポートする必要があります。

combinations_with_replacement()は何をする関数かというと、『重複組み合わせ』を列挙する関数です。普通の『組み合わせ』combinations()との違いは、同じ要素を複数回選ぶのを許するところです。

入力は『イテラブル』な要素と、要素を取り出す回数（列の長さ）の2つです。『イテラブル』は英語で"iterable"で、「反復できる」という意味です。forループのinの後に使えるもののことで、「リスト」、「タプル」、「range()」、「文字列」などがあります。

出力は、入力の条件でできるすべての『重複組合せ』です。出力はそれぞれ、「入力された要素の順番」で並べられてでてきます。これは重要なことです。

forループとprint()で、何が作られるのか見てみましょう。要素をrange(1,4)、つまり(1,2,3)の3つで、長さは2としてみます。

# 長いので、comb_rplcという名前でインポートします
from itertools import combinations_with_replacement as comb_rplc
for seq in comb_rplc(range(1, 4), 2):
    print(seq)

(1, 1)
(1, 2)
(1, 3)
(2, 2)
(2, 3)
(3, 3)

（1, 1)や(2, 2)のように、同じ要素が複数回出てくるものも含んでいます。入力は1,2,3の順番なので、(2, 1)や(3, 2)のように、前の数字より小さい数字が次にくることはありません。

普通のcombinationsと比べてみる

普通の組み合わせcombinations()と比べてみましょう。

from itertools import combinations
for seq in combinations(range(1, 4), 2):
    print(seq)

(1, 2)
(1, 3)
(2, 3)

たしかに、同じ要素を2回選んでいる(1, 1)や(2, 2)は出てきていません。

入力の順番を変えてみる

入力を"CBA"にしてみます。forと同じように、"C"、"B"、"A"の3要素とみなされます。

for seq in comb_rplc("CBA", 2):
    print(seq)

('C', 'C')
('C', 'B')
('C', 'A')
('B', 'B')
('B', 'A')
('A', 'A')

C,B,Aの順番で入力したので、出力もその順番に並べられて出てきます。ABCの順に出てくるわけではないことに注意しましょう。

この問題の数列をどう作るか？

この問題で作る数列の条件は、次の数字が前の数字より小さくないことです。入力を昇順にすれば、出力も昇順になるので、この条件は勝手に満たしてくれます。

数列の長さは $N$ で、数字の上限が $M$ なので、combinations_with_replacement(range(1, m + 1), n)とすれば、すべてのありえる数列を列挙できます。

コード1

from itertools import combinations_with_replacement as comb_rplcのように省略名をつけてインポートすれば、コードがごちゃごちゃしなくて済みます。

from itertools import combinations_with_replacement as comb_rplc

n, m, q = list(map(int, input().split()))
# reqは[[a1,b1,c1,d1],[a2,b2,c2,d2]……]が入ったリストのリストです
req = [list(map(int, input().split())) for _ in range(q)]

ans = 0
# seqは長さnのタプルです
for seq in comb_rplc(range(1, m + 1), n):
    score = 0
    for a, b, c, d in req:
        # 問題文に書いてある数列のk番目は、インデックスだとk-1になるので注意
        if seq[b - 1] - seq[a - 1] == c:
            score += d
    ans = max(ans, score)

print(ans)

余談1："with replacement"ってなんだ

"with replacement"は、「元に戻す」という意味です。"replace"は「交換する」、「置き換える」という意味のほうがよく使われますが、そちらではありません。

数字の書いたボールが袋に入っているとします。"without replacement"な普通の『組み合わせ』では、取り出したボールを袋に戻しません。"with replacement"な『重複組み合わせ』では、取り出したボールの数をメモしたあと、袋に戻します。

こういうふうに意味がわかると、この関数を思い出しやすくなるかもしれません。

方法2 キュー再帰で作る方法

この問題はcombinations_with_replacement()で数列を作ることができましたが、条件がもっと複雑になると、自力で実装する必要が出てきます。

「ある条件を満たす文字列や数列を全て作りたい」とき、汎用性の高い方法に『キュー再帰』というものがあります。

これはその名の通り、キュー（queue）というデータ構造（配列）を使います。キューはどういうものかというと、「先に入れたものを、先に取り出す」（FIFO: First In, First OUT）配列です。アルゴリズムの勉強をすると出てくる、メジャーなやつです。

キューからまだ作りかけの文字列を取り出して、そこに1文字追加した文字列たちを、キューに追加することを繰り返します。

すると、最終的に作りたい数列が全て生成されます。

Pythonでキューを使うには、"deque"を使う

Pythonでキューを使うには、collectionsモジュールのdeque()をインポートする必要があります。

"deque"は"double-ended-queue"の頭文字をとったもので、『両端キュー』というデータ型です。発音は「デック」です。

dequeは先頭、末尾のどちらからでも、要素を取り出したり、追加したりすることができます。つまり、スタックとキューの上位互換です。

dequeのメソッド

dequeに要素を追加、取り出すメソッドは2つずつあります。

append()：右（末尾）に要素を追加する
appendleft()：左（先頭）に要素を追加する
pop()右（末尾）の要素を取り出す
popleft()：左（先頭）の要素を取り出す

キューの「先に入れたものを、先に取り出す」とは、「入れるときは右から」「取り出すときは左から」と言い換えられます。

つまり、入れるときはappend()を使って、取り出すときはpopleft()を使えば、それだけでキューが実現できます。

スタックの場合は「後に入れたものを、後に取り出す」なので、append()で取り出すときはpop()です。

どんなコード？

キュー再帰を疑似コード風に書いたものをみて、どういうことをするのか見てみます。

from collections import deque
que = deque()

que.append(最初の状態)  # 最初の状態を入れないと、while queを素通りします

while que:
    今の状態 = que.popleft()  # 先頭から取り出します
    # なにか処理をする
    if 条件を満たす:
        # なにか処理をする
    else:
        次の状態 = 今の状態 + α # forで「次の状態」を複数作る場合もある
        que.append(次の状態)  # 末尾に追加します

while queとは、queが空でないならTrue、空になるとFalseと判定されるので、queの中身がなくなるまでループするという意味です。

条件を満たす場合、キューに追加せずに何らかの処理をしないと、無限ループになります。

こうすると、キュー再起ができます。

コード2

from collections import deque


# 数列の点数を計算する関数
def calc(seq):
    score = 0
    for a, b, c, d in req:
        if seq[b - 1] - seq[a - 1] == c:
            score += d

    return score


n, m, q = list(map(int, input().split()))
req = [list(map(int, input().split())) for _ in range(q)]

ans = 0

que = deque()

# 数列の1番目、[1]～[m]までキューに追加しますが、
# 実は、この問題は数列の最初が[1]の場合だけを考えても解けます。
for i in range(1, m + 1):
    que.append([i])

while que:
    seq = que.popleft()

    if len(seq) == n:
        # 長さがnになったので、得点を計算します
        score = calc(seq)
        ans = max(ans, score)
    else:
        # 次に追加する数字は、下限が今の数列の一番後ろの数字、上限がmです
        for i in range(seq[-1], m + 1):
            seq_next = seq + [i]
            que.append(seq_next)

print(ans)

キューの中身の変化を見てみる

m=3、n=3のとき、キューの中身がどう変化していくのか、書いてみます。

[1],[2],[3]（初期状態）
[2],[3],[1,1],[1,2],[1,3]
[3],[1,1],[1,2],[1,3],[2,2],[2,3]
[1,1],[1,2],[1,3],[2,2],[2,3],[3,3]
[1,2],[1,3],[2,2],[2,3],[3,3],[1,1,1],[1,1,2],[1,1,3]
[1,3],[2,2],[2,3],[3,3],[1,1,1],[1,1,2],[1,1,3],[1,2,2],[1,2,3]
……（中略）
[2,3,3],[3,3,3]
[3,3,3]
なし（終了）

一番左にあるものを取り出して、それに1個数字を付け加えたものを右に追加していっていますね。このようにして、すべての数列を作り出すことができます。

長さが3は数列が完成したということなので、点数を計算するだけで、新しい状態をキューに追加することはありません。そのため、最終的にキューの中身は空になるので、無限ループに陥ることはありません。

余談1：スタック再帰でも解けます

この問題の場合、「右から入れて左から取り出す」キューでなく、「右から入れて右から取り出す」スタックでも解けます。

ついでに、スタックはdequeでなく、普通のリストでもできます。append()メソッドとpop()メソッドは普通のリストにもあるからです。

しかし、基本的にはdeque()をキューとして扱って解くことをおすすめします。3つほど利点があります。

要素数が増えるとdequeのほうが動作が速い
キューを使うと数列や文字列が『辞書順』で出てくることがたまに役立つ
文字列・数列列挙以外でもdequeを使う機会はあるので、慣れておくと有利

余談2：キュー再帰は『幅優先探索』で、スタック再帰は『深さ優先探索』

キュー再帰は『幅優先探索』（BFS）そのものです。スタック再帰にすると『深さ優先探索』（DFS）になります。

文字列や数列をすべて作るだけならどちらでもできますが、出てくる順番が違います。

方法3 深さ優先探索（DFS）で作る方法

最後は、公式の解説PDFに書いてある、深さ優先探索を使う方法です。『再帰関数』を使います。

キュー再帰やスタック再帰よりもさらにできることの幅が広がりますが、文字列や数列を作るだけならキュー再帰だけで十分です。

深さ優先探索のメリットは

『行きがけ処理』『帰りがけ処理』を実装しやすい
できるとアルゴリズムを書いてる感がある

デメリットは

再帰関数は動作を理解しづらい（しようとしないほうがいいです。感覚と慣れです）
Pythonだと再帰回数の上限に引っかかってREになりがち

コード3

dfs関数の引数は、今の数列seq（タプル）です。次の数字の下限は、今の数列の一番後ろなので、seq[-1]です。

例えば、数列1,1に3を付け加えて、1,1,3になったとします。次に付け加える数字は3以上でなければいけないので、(次の数字の下限)は今付け加えた3です。

今の数列の長さがn未満の場合は、(数字の下限)~(数字の上限m)を付け加えた新しい数列をそれぞれすべて、再びdfs(今つくった数列、今つくった数列の長さ、次の数字の最小値)をします。

今の数列が1,1,3で、上限$M = 4$ならば、
1,1,3,3
1,1,3,4
をdfsに渡すということです。

長さがnになって完成したら、得点を計算します。得点が出たら、ans = max(ans, score)として最大値を更新します。これは普通の問題と同じです。

すべてが終わったら、本体のコードにすべての最大値が返ってくるので、これをprintして終わりです。

def dfs(seq):
    # 返り値は、すべての数列の得点の最大値 ans です。
    ans = 0
    if len(seq) == n:
        # 数列が完成したので、得点を計算します
        score_ret = 0
        for a, b, c, d in req:
            if seq[b-1] - seq[a-1] == c:
                score_ret += d
        return score_ret  # この数列の得点を返します
    else:
        # まだ数列が完成していません
        for i in range(seq[-1], m + 1):
            seq_next = seq + (i,)  # 長さ1のタプル(i,)を連結します
            score = dfs(seq_next)  # seq_nextから派生するすベての数列の中での、得点の最大値が返ってきます
            ans = max(ans, score)  # 最大の得点を更新します

    # 得点の最大値を返します
    return ans


n, m, q = list(map(int, input().split()))
# reqは[[a1,b1,c1,d1],[a2,b2,c2,d2]……]が入ったリストのリストです
req = [list(map(int, input().split())) for _ in range(q)]

# 最終的に答えが返ってくるようにします。処理はすべてdfsメソッドでやってもらいます。
# リストだとどこかで間違えて値を書き換えそうで怖いので、タプルにしておきます
# 1番目が1の場合以外は考えなくていいので、(1,)だけやります
ans = 0
score = dfs((1,))
ans = max(ans, score)
print(ans)

1番目が1の場合以外考えなくていい理由

何度か数列の1番目が1の場合以外考えなくてもいいと書きました。コード3の深さ優先探索では、実際に1始まりだけを全探索しています。

なぜそうなるかというと、「b番目の要素とa番目の『差』」が重要であって、数字そのものはなんでもいいからです。

たとえば、2,3,3,4は1,2,2,3と差がそれぞれ同じですし、3,3,3,3は1,1,1,1と同じです。このように、1始まり以外の数列には、1始まりで等価な数列が必ずあります。

別に全部やってもいいのですが、気づくと実装がちょっと楽になるかもしれません。ついでに実行時間も半分くらいになりますが、1始まり以外を含めて全探索してもTLEにならないのでどうでもいいです。

類題

類題をいくつか紹介します。一番上の問題は慣れるのにちょうどいい難易度です。後ろの2つはやや難しいD問題レベルの難易度ですが、それでもこの問題よりは簡単です。

茶レベル： ABC029 C - Brute-force Attack
緑レベル： ABC161 D - Lunlun Number
緑レベル：パナソニックプログラミングコンテスト2020 D - String Equivalence

【AtCoder解説】PythonでABC165のC問題『Many Requirements』を制する！