AtCoder Beginner Contest 095に事後参加してみた
この回はD問題がけっこう面白かった。ちょっとコードも解説も多めかもしれない・・・
第1問
700円のラーメンに味玉、チャーシュー、ネギのトッピング(各100円)するか否かを示す'o'か'x'かが三つ並んだ文字列を与えられて、必要な金額を計算する問題。
print(700 + 100*input().count('o'))
そのまんま。
第2問
ドーナッツの種類がN個あり、種類ごとに必要な粉の量mがあり、全種類のドーナッツを少なくとも一つは作る必要があるという制約のもと、粉の量Xで最大いくつのドーナッツを作ることができるか、という問題。
n, x = map(int, input().split()) ms = [int(input()) for _ in range(n)] print(n + (x - sum(ms))//min(ms))
全種類作るということでN個のドーナッツをmの和の量の粉で作って、残った粉でもっともコストの低いドーナッツを最大限作った時の数を返している。
第3問
二種類のピザA、Bと、その二種類が半分ずつ合体したピザABが注文できる。AやBがほしい場合、AやBをそのまま買うか、ABを二枚買って半分にわけて一枚のAと一枚のBにする方法がある。AをX枚、Bを Y枚ほしい場合に一番安く買える金額を求める。
決定するべき分岐点は二つある。
まず「AとBのどちらも必要な数(つまりXとYの最小値)を個別で買って揃えるか、ABを買って揃えるか」でそれは(a+b)と(c*2)のどちらが安いかで決まる。
次に「残ったAかBの必要な数を個別で買うか、ABを買って余った部分は捨てるか」でそれはaと(c2)、あるいはbと(c2)のどちらが安いかで決まる。
ちなみにZがXとYの最小値として、(X-Z)と(Y-Z)のいずれかは0。
というロジックをコードにするとこうなる:
a, b, c, x, y = map(int, input().split()) z = min(x, y) total = z * min(a + b, 2 * c) total += (x - z) * min(a, 2 * c) total += (y - z) * min(b, 2 * c) print(total)
第4問
円形のカウンターに置いてある寿司に、スタート地点からの時計回りでの距離(=到達までのカロリー消費量)と寿司を食べた時に得られるカロリーの二つの値がある。時計回り、反時計回りに動き回って得られるカロリー値を最大化する問題。
時計回り、反時計回りの計算をなるべく対称にするため、「スタート地点からの時計回りでの距離のリスト」を「点と点の間の距離のリスト」に変換しておく。
ds = [j - i for i, j in zip(chain([0], xs), chain(xs, [c]))]
xsが各点の起点からの距離、dsが起点を最初と最後に入れた上での、各点どうしの距離のリスト。cはカウンターの全長。
一方向しか動かない縛り
手始めに問題を単純化して、どちらかの方向にしか進めない場合を考えてみる。
時計回りに動いた場合、各点まで到達することで得られるカロリーはその点までの寿司のカロリーの合計からその点までの距離の合計を引いたものだ。
pythonでaccumulate関数が「リストの添字nまでの暫定合計」を返してくれるので、それを利用するとこう書ける:
clockwise = chain([0], (v - d for v, d in zip(accumulate(vs), accumulate(ds))))
(vsは各点に置かれている寿司のカロリー値)
とりあえず「全く動かない場合に得られるカロリー」ということでシーケンスの先頭に0をいれておいた。
これを時計回り・反時計回りでやって最大値を取れば、一方向に進むだけの場合の最大カロリーがわかる。
clockwise = chain([0], (v - d for v, d in zip(accumulate(vs), accumulate(ds)))) anticlockwise = chain([0], (v - d for v, d in zip(accumulate(vs[::-1]), accumulate(ds[::-1])))) maximum = max(chain(clockwise, anticlockwise))
chainでイテレータを二つつなげて、maxでその最大値をとっている。
方向転換も可能な場合
上記を踏まえて、方向転換について考えてみる。
少し考えると、二回以上方向転換をする意味がないことがわかる。
二回方向転換しても二度目の方向転換の後は「すでに寿司をとってしまった距離を再度歩いている」だけなので、「時計回り・反時計回り・時計回り」といった進みかたをするより「反時計回り・時計回り」と進んだほうが同じ寿司をより少ない距離でとれる。
なので方向転換を考慮に入れなければいけないのは一回のみ。
さらに「方向転換しない」というのは「まず一方向に0距離進んでから逆方向に転換する」という「一回方向転換」の特殊ケースとして扱える。
ということでこの問題は「まず時計回りにすすんでから方向転換して反時計回り」「まず反時計回りにすすんでから方向転換して時計回り」の二つのケースを調べるだけで済む。
m点まで時計回りで進んで、起点に戻って、反時計回りで起点からm+1点まででカロリーが最大化される点に進む
あるいは
m点まで反時計回りで進んで、起点に戻って、時計回りで起点からm-1点まででカロリーが最大化される点に進む
の二ケースだ。
時計回りから反時計回りに転換するケースを考えよう。
まず時計回りにmの地点まで進む場合、得られるカロリーはmまでの寿司のカロリーの合計。使われるカロリーはmまでの距離の二倍だ。なぜ二倍かというと、「起点まで戻る」というコストもかかるから。
clockwise1 = chain([0], (v - d*2 for v, d in zip(accumulate(vs), accumulate(ds))))
いったん起点に戻ってから反時計回りに進む場合の得られるカロリーは前出の一方向のケース:
anticlockwise1 = chain([0], (v - d for v, d in zip(accumulate(vs[::-1]), accumulate(ds[::-1]))))
では時計回りにmの地点まで進んでから反時計回りに進んで得られる最大のカロリーは?となると「反時計回りで起点からm+1点まででカロリーが最大化される点に進む」必要が出てくる。anticlockwise1の最大値を求めるだけではだめだ。起点nからm+1までの[n, n-1, n-2, n-3 ..., m+1]における暫定最大値を計算する必要がある。
というわけでまた出てくるのがaccumulate。デフォルトでは暫定合計をリスト化するが、関数fを渡すことで暫定reduce(xs, f)をリスト化してくれる。なのでmax関数を渡すことで、[max(xs[:1]), max(xs[:2]), max(xs[:3]) .., max(xs[:n])]というリストをO(len(xs))で作成する。
accumulate(xs, func=max)を使って「m地点まで時計回りで進んでから起点に戻って反時計回りに進んだ場合のカロリー最大値」のリストを作成するpythonコード:
clock_anticlock = (x + y for x, y in zip(clockwise1, reversed(list(accumulate(anticlockwise1, func=max)))))
これを反時計回りでもやれば問題解決。
ということで全部のコード:(コード部分横スクロールできます)
from itertools import accumulate, chain n, c = map(int, input().split()) xs, vs = zip(*(map(int, input().split()) for _ in range(n))) ds = [j - i for i, j in zip(chain([0], xs), chain(xs, [c]))] clockwise1 = chain([0], (v - d*2 for v, d in zip(accumulate(vs), accumulate(ds)))) anticlockwise1 = chain([0], (v - d for v, d in zip(accumulate(vs[::-1]), accumulate(ds[::-1])))) clock_anticlock = (x + y for x, y in zip(clockwise1, reversed(list(accumulate(anticlockwise1, func=max))))) anticlockwise2 = chain([0], (v - d*2 for v, d in zip(accumulate(vs[::-1]), accumulate(ds[::-1])))) clockwise2 = chain([0], (v - d for v, d in zip(accumulate(vs), accumulate(ds)))) anticlock_clock = (x + y for x, y in zip(anticlockwise2, reversed(list(accumulate(clockwise2, func=max))))) print(max(chain(clock_anticlock, anticlock_clock)))
とりあえずO(n)。比較的複雑そうなロジックがzip, accumulate, chainを使って効率良く記述できるのは嬉しい。
感想
コードを書くよりブログを書くほうが何倍も難しかった・・・ 言語化って大変だな。