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