Python で解く AtCoder ABC128

2019年5月26日に行われたAtCoder Beginner Contest 128のPython3 での実装例です

A - Apple Pie

使用テクニック： 切り捨て除算 (//)

B - Guidebook

使用テクニック： 条件付きソート (sorted)

sorted 関数の引数 key にラムダ関数を使用してレストラン名、ポイントを、この順番で渡します。そのソート後の順番で、入力されたレストランの番号を出力します。

C - Switches

使用テクニック： bit全探索

方針

全てのスイッチ状態を調べ、全ての電灯が点灯している状態を数えます。
最大でスイッチの数が $10$ 個、スイッチの状態はon/off の2通りであるため、スイッチの取り得る状態の数は最大で $2 ^ {10} = 1024$ 通りです。また、電球の数も最大で $10$ 個、電球につながっているスイッチの数の最大も $10$ 個です。そのため、全てのスイッチの状態において、全ての電球の、全てのスイッチの状態を調べるときの計算量は最大で $10 \times 10 \times 2 ^ {10} \approx 10 ^ 6$ となり、十分計算できる範囲に収まります。
ちなみにスイッチの数が $N$ , 電球の数が $M$ であるときの計算量を一般化すると $O(N M \cdot 2 ^ N)$ となります。

実装

スイッチの状態を $0$ ～ $2^N - 1$ の数値で表し、各数値を 2進数表記した際に $1$ となっている（bit が立っている)位置のスイッチが on になっているとします。
各電球に対して、繋がっているスイッチの番号の位置のbitが立っているかを確認することで、各電球の on　になっているスイッチの数を調べます。（17-20行目）
そして、onになっているスイッチの数から、その電球が点灯しているかを調べ（22-25行目）、全ての電球が点灯していれば答えに 1 を足します。（27-28行目）

D - equeue

使用テクニック： Priority queue

方針

問題を分割して考えます。まず、既に $n$ 個の宝石を持っており、ここから $m$ 個以内の宝石を捨てて宝石の価値の総和を最大化する問題を考えます。
この場合、価値がマイナスである宝石の中で、最も価値の低い宝石から捨てていくことで、価値の総和を最大化できます。
また、捨てる宝石の数が $m$ 個に達しなくても、価値がマイナスの宝石が無くなればそこで捨てるのを止めることが最適です。
あとは、取り得る全ての宝石の持ち方において、宝石の価値を最大化し、その最大値が答えとなります。

具体的には、dequeue $D$ の左から宝石を取る数と、右から取る数としてあり得るパターンの全通りを試します。
左からの宝石の取り方は、 $\min(N, K)$ 通りあります。左から取った宝石の数を $r$ とすると、右からの宝石の取り方は、右で既に取った宝石を取ることはできないため、 $\min(N, K) - r$ 通りあります。
そして、左からとった宝石の数を $l$ とすると、宝石を捨てることができる回数は $K - (r + l)$ となります。
捨てる宝石の取得には Priority queue (heapq) を使用すれば良く、 $n$ 個の宝石の中から捨てる宝石を取得する計算量は $\log n$ です。

以上より計算量を概算します。
$\min (N, K) = R$ とすると、左からの宝石の取り方が $R$ 通り、右からの宝石の取り方も $R$ 通りあります。
宝石を捨てる回数は最大 $K$ 回、捨てる宝石の取得の計算量が最大 $\log R$ です。
以上より、この実装の計算量は $O(R^2 K \log R)$ となります。

実装

左からの宝石の取り方と右からの宝石の取り方を全通り試すため、二重 for-loop を回します。
そして、操作数が $K$ に達するか、取得した宝石から全ての価値がマイナスの宝石を捨てるまで while-loop を回します（18-24行目）。

E - Roadwork

使用テクニック： 二分探索

方針

座標 $0$ に近い通行止め地点から順番に、そこで足を止める人を決めていきます。
全ての人が座標 $0$ から速度 $1$ で正方向に歩くため、地点 $X _ i$ で時刻 $(S _ i - 0.5, T _ i - 0.5)$ の間に通行止めになっていた場合、時刻 $(S _ i - X _ i - 0.5, T _ i -X _ i - 0.5)$ の間に地点 $0$ を出発した人はこの通行止めに引っ掛かります。出発時刻はすべて整数値なので、通行止めに引っ掛かる人の出発時刻は半開区間で $[S _ i - X _ i, T _ i - X _ i)$ と表すことができます。

よって、座標 $0$ に近い $X _ i$ から順に $[S _ i - X _ i, T _ i - X _ i)$ の間に出発している人を探すことで、各人の歩行を止める地点が決まります。一度も $[S _ i - X _ i, T _ i - X _ i)$ に含まれなかった人は無限に歩き続けることができます。
時刻 $[S _ I - X _ i, T _ i - X _ i)$ に含まれる人の探索には二分探索を使用します。通行止めの開始時刻と終了時刻のそれぞれで二分探索を行い、その間に出発していて、まだ止まっていない人が $X _ i$ で歩くのを止める人です。

この実装の計算量を概算します。
まず $N$ 個の通行止め地点の昇順ソートが $O(N \log N)$ 、1回の二分探索の計算量が $O(\log Q)$ であり各通行止め地点につき $2$ 回の二分探索を行うため、二分探索全体で $O(2N \log Q)$ となります。そして $Q$ 人の止まる点を判断するために $O(Q)$ の計算量が必要となります。
合計でこの実装の計算量は $O(N \log N + 2N \log Q + Q)$ となります。

実装

提出 #5654623 - AtCoder Beginner Contest 128 を参考にさせてもらいました。

通行止め地点をリスト closure に保持し、ソートしておきます。（9-14行目）
また、 $Q$ 人の出発時刻をリスト D に保持します。制約から D は既に昇順でソートされています。（16行目）

そして、各通行止め地点で止まる人を探索します。（20-30行目）
二分探索 bisect_left を使用して地点 $X _ i$ で止まる人は、リスト D のインデックス left - right の間に位置している人になります。ただし、この間に位置していても既に止まっている可能性があります。

そこで、各人の止まる地点をメモするリスト ans と、その人が既に止まっているかをメモしておくリスト skip の2つのリストを用意しておきます。skip は、まだその人が止まっていなければ $-1$ を入れておきます。既に止まっている場合は、同じ通行止め地点で止まった右端の人のインデックスを入れておきます。
この skip リストのおかげで、同じ人に対して何度も止まっているかどうかの判定をする必要がなくなり、 $Q$ 人の止まる地点の判断に必要な計算量が $O(Q)$ で済みます。

F - Frog Jump

使用テクニック： 動的計画法 (DP)

方針

問題文の言い換え
まず蓮が無い地点に飛んでしまった場合の減点が非常に大きいため、同じ蓮に飛んだ時に最終得点が最大になることはあり得ません。よって、同じ蓮に飛ぶことなく最後に $N-1$ 番目の蓮に飛ぶルートのみを考えます。
あとは全てのあり得るルートを求め、最大の最終得点を求めます。

問題文より、位置の遷移は、
$0, A, A-B, 2A-B, 2A-2B, \cdots, (n+1)A-nB$
となります。また、最後には $N-1$ 番目の蓮に飛ばなければならないため
$(n+1)A - nB = N-1$ 式(1)
である必要があります。

ここで、 $A-B = C$ と置くと、位置の遷移は次のように書き換えることができます。
$0, A, C, \cdots, A + kC, kC, \cdots, nC, A + nC$ $(0 \leqq k \leqq n)$
この数列は、さらに次のように変形できます。
$0, C, \cdots, kC, \cdots, nC, A, \cdots, A+kC, \cdots, A+nC$

また、 $A-B=C$ より式(1) は、
$A + nC = N-1$
となります。

以上より、蓮の飛び方の全てのルートを試すということは、 $(A, C)$ の全ての組合せを試すということに言い換えることができます。

新しい制約
ここで、 $(A, C)$ には以下の制約があります。

$A$ は $1$ 以上の整数
$C$ は $A$ 以下の整数
$C = A - B$ かつ $B$ は正の整数 $(B \geqq 0)$ であるため、 $C \leqq A$
$C$ は $\frac{N}{2}$ 未満の整数
$A + nC = N-1$ と $A \geqq C$ より、 $C + nC \leqq N-1$ すなわち $C \leqq \frac{N-1}{n+1}$ となる
$n=0$ は $0$ 番目から $N-1$ 番目の蓮へ直接飛ぶパターンであるため $C$ は任意の数でよいが、この場合の得点は必ず $0$ であるため考慮しない
よって $n=1$ の時に $C$ は最大となる $(C \leqq \frac{N-1}{2})$

動的計画法による最終得点の求め方
愚直に全ての $(A, C)$ の組合せを列挙して最終得点を求めると TLE になってしまいます。そこで、 $C$ を固定してジャンプの数を徐々に増やしていった時の最終得点の変化を考えます。 $k$ 回ジャンプした時は以下の番号の蓮に乗っています。
$0, C, \cdots, kC, A,_k \cdots, A _ k+kC$ 制約より、 $A _ k+kC=N-1$

同様に $k+1$ 回ジャンプする時は以下の番号の蓮に乗ります。
$0, C, \cdots, kC, (k+1)C, A _ {k+1}, \cdots, A _ {k+1} + (k+1)C$
また、制約より $A _ {k+1} + (k+1)C = N - 1$ を満たす必要がありますが、この式より $A _ k$ は次の漸化式で表すことができます。
$A _ {k+1} = N - 1 - kC - C = A _ k - C$

よって、 $k+1$ 回ジャンプするときに乗る蓮の番号は、 $k$ 回ジャンプするときの蓮に $(k+1)C, A _ k - C$ 番目の2つの蓮を追加したものになります。
よって、 $C$ が固定されており、 $k$ 回ジャンプした時の最終得点を $S _ {C, k}$ とし、 $m$ 番目の蓮の得点を $P _ m$ としてあらわすと、 $S _ {C, k}$ は次の漸化式で表すことができます。
$S _ {C, k} = S _ {C, k-1} + P _ {A_k} + P _ {kC}$ $(S _ {C, 0} = 0)$

$k$ の範囲
$A+nC=N-1$ かつ、 $A \geqq 1$ より、 $k \lt \frac{N-1}{C}$ となります。 $k$ はジャンプする回数であるため、その最大値は $\frac{N-1}{C}$ 未満の整数となります。

ただし、 $C$ が $N-1$ を割り切る場合は注意が必要です。
この時、 $C$ は $N-1$ の約数です。加えて、 $A = N-1 - kC$ であるため、 $A$ も $N-1$ の約数となります。そのため、 $kC \geqq A$ のとき $kC$ 番目の蓮に飛ぼうとしても、その蓮は既に使用済みになっています。
よって、 $C$ が $N-1$ を割り切る場合は $k \lt \frac{A}{C}$ である必要があります。

また、 $A$ は $C$ と $k$ が決まれは $A = N-1 - kC$ より一意に求めることができますが、 $C \leqq A$ の制約にも注意する必要があります。

計算量
この実装では $(A, C)$ の全ての組合せを探索しますが、 $A$ は $C$ とジャンプの回数 $k$ によって一意に定まるため $(C, k)$ の全ての組合せを探索すると言い換えることができます。 $C$ の範囲は $[1, \frac{N}{2}]$ 、 $k$ の範囲は $[1, \frac{N-1}{C}]$ であるため、 $(A,C)$ 全列挙の計算量は $O(N \log N)$ となります。
また、各組合せの最終得点は動的計画法により $O(1)$ で求めることができるため、この実装全体の計算量は $O(N \log N)$ となります。