Python で解く AtCoder ABC129

2019年6月9日に行われたAtCoder Beginner Contest 129のPython3 での実装例です

A - Airplane

使用テクニック： ソート (sort)

入力を昇順でソートし、値の小さい2つを足したものを答えます。

B - Balance

使用テクニック： メモ化

グループの分割位置を全て試し、最小となる答えをメモしておきます。

C - Typical Stairs

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

方針

階段の上り方として1段先と2段先のどちらにも進めるため、 $i$ 段目の階段に到達する移動方法は、 $i-1$ 段目に到達する方法と $i-2$ 段目に到達する方法の和となります。ただし、 $i$ 段目の階段が壊れている場合、その階段を使うことはできないためこの階段に到達する方法は $0$ となります。

よって、壊れている階段の情報をリストに保存しておき、 $i$ 段目の階段が壊れているかどうかで場合分け行います。

　実装

注意点として、 $1$ 段目の階段に上るときに $2$ 段前の階段到達方法を参照しようとしても $-1$ 段目となってしまい参照範囲外となってしまいます。そこで、初めから $1$ 段目にいると考え、 $N+1$ 段目に到達する方法の数を求めるものとします。この場合、壊れている階段の段数が入力から $1$ 段上になることに注意してください（10行目）。

D - Lamp

使用テクニック： 動的計画法(DP), numpy.ndarray

方針

愚直に各マスについて行と列方向の連続マスを数えると計算量が $O(HW(H+W)$ となり、時間内に計算を終わらせることができません。

そこで $H \times W$ のグリッドを4個用意し、それぞれ上・下・左・右の4方向の端から障害物のないマスが何マス連続しているかを行または列ごとに数えます。障害物があった場合、そこでカウントをリセットしてまた１から数えなおします。

具体例として、次のような $3 \times 4$ グリッドを考えます。この場合、座標 $(1, 3)$ が照らすことのできるマスを最大化し、その答えは $6$ になります。

$\begin{matrix} . &. &. &. \\ . &\# &. &\# \\ \# &. &. &. \end{matrix}$

まず、このグリッドを上・下、左・右の端からカウントを行うと次のようになります。

$top = \begin{pmatrix} 1& 1 & 1 & 1 \\ 2 & 0 & 2 & 0 \\ 0 & 1 & 3 & 1 \end{pmatrix}, bottom = \begin{pmatrix} 2& 1 & 3 & 1 \\ 1 & 0 & 2 & 0 \\ 0 & 1 & 1 & 1 \end{pmatrix}$

$left = \begin{pmatrix} 1 & 2 & 3 & 4 \\ 1 & 0 & 1 & 0 \\ 0 & 1 & 2 & 3 \end{pmatrix}, right = \begin{pmatrix} 4& 3 & 2 & 1 \\ 1 & 0 & 1 & 0 \\ 0 & 3 & 2 & 1 \end{pmatrix}$

そして、上・下のグリッドを足し合わせ、障害物のないマスから $1$ を引くと列方向の連続マス数を求めることができます。同様の操作を左・右のグリッドに行うと、行方向の連続マス数を求めることができます。

$\begin{align} C &= top + bottom - empty \\ &= \begin{pmatrix} 1& 1 & 1 & 1 \\ 2 & 0 & 2 & 0 \\ 0 & 1 & 3 & 1 \end{pmatrix} + \begin{pmatrix} 2& 1 & 3 & 1 \\ 1 & 0 & 2 & 0 \\ 0 & 1 & 1 & 1 \end{pmatrix} - \begin{pmatrix} 1& 1 & 1 & 1 \\ 1 & 0 & 1 & 0 \\ 0 & 1 & 1 & 1 \end{pmatrix} \\ &= \begin{pmatrix} 2& 1 & 3 & 1 \\ 2 & 0 & 3 & 0 \\ 0 & 1 & 3 & 1 \end{pmatrix} \end{align}$

$\begin{align} R &= top + bottom - empty \\ &= \begin{pmatrix} 1& 2 & 3 & 4 \\ 1 & 0 & 1 & 0 \\ 0 & 1 & 2 & 3 \end{pmatrix} + \begin{pmatrix} 4& 3 & 2 & 1 \\ 1 & 0 & 1 & 0 \\ 0 & 3 & 2 & 1 \end{pmatrix} - \begin{pmatrix} 1& 1 & 1 & 1 \\ 1 & 0 & 1 & 0 \\ 0 & 1 & 1 & 1 \end{pmatrix} \\ &= \begin{pmatrix} 4 & 4 & 4 & 4 \\ 1 & 0 & 1 & 0 \\ 0 & 3 & 3 & 3 \end{pmatrix} \end{align}$

最後に行方向と列方向の連続マス数を足して障害物のないマスから $1$ を引くと、各マスの行・列方向の連続マス数を求めることができます。この行列の最大値が求める答えとなります。

$\begin{align} RC &= R + C - empty \\ &= \begin{pmatrix} 2 & 1 & 3 & 1 \\ 2 & 0 & 3 & 0 \\ 0 & 1 & 3 & 1 \end{pmatrix} + \begin{pmatrix} 4 & 4 & 4 & 4 \\ 1 & 0 & 1 & 0 \\ 0 & 3 & 3 & 3 \end{pmatrix} - \begin{pmatrix} 1 & 1 & 1 & 1 \\ 1 & 0 & 1 & 0 \\ 0 & 1 & 1 & 1 \end{pmatrix} \\ &= \begin{pmatrix} 5 & 4 & 6 & 4 \\ 2 & 0 & 3 & 0 \\ 0 & 3 & 5 & 3 \end{pmatrix} \end{align}$

実装

入力されるグリッドを障害物のないマスを $1$ 、障害物の有るマスを $0$ とした行列で保持します。(8-12 行目)

上・下・左・右の連続マスを数えるための4つの行列を用意します。この時、初めの行または列を数えるために、1行または1列を追加しておきます。(13-16 行目)

4つの行列について、各行・列の値は前の行・列に $1$ を加えた値になりますが、障害物がある場合はカウントをリセットするために $0$ を掛けます。(18-28 行目)

E - Sum Equals Xor

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

方針

非負整数 $a, b$ が $a + b = a \, \mathrm{xor} \, b$ となるためには、2進数表記した際の $i$ 桁目を $a _ i, b _ i$ とした際に $(a _ i, b _ i)$ の組合せが $(0, 0), (1, 0), (0, 1)$ のどれかである必要があります。

また、 $a + b \leqq L$ となる $(a, b)$ の組合せを考えます。
$L$ の上位桁から見ていき、 $L$ の上から $i$ 桁目を $L _ i$ とします。 $L _ j = 1$ であったときに、 $i \lt j$ では $a _ i + b _ i = L _ i$ であり、 $a _ j + b _ j = 0$ となるような $(a, b)$ の組合せであれば、 $j$ 以降の桁が何であっても $a + b \leqq L$ となり、条件を満たすことが確定します。一方、 $a _ j + b _ j = 1$ となるような $(a, b)$ の組合せであれば、 $j$ 桁以降の選択によって $a + b \gt L$ となる可能性があります。

以上を踏まえて、 $L$ の $i$ 桁目まで条件を満たすことが確定した組合せの数を $dp0 _ i$ , 未確定の組合せの数を $dp1 _ i$ とすると以下の漸化式が成立します。

$L _ {i+1} = 0$ の場合、
条件を満たしていることが確定している組み合わせでは $i+1$ 桁目として $(a _ i, b _ i) = (0, 0), (1, 0), (0, 1)$ のどれを選んでも良いため $dp0 _ {i+1} = dp0 _ i \times 3$ となります。
一方、条件を満たすことが未確定である場合、 $i+1$ 桁目としては $(a _ i, b _ i) = (0, 0)$ のパターンでのみ $a + b \leqq L$ を満たす可能性が残ります。よって、 $dp1 _ {i+1} = dp1 _ i \times 1$ となります。

$L _ {i+1} = 1$ の場合、
$dp1 _ i$ に含まれる組合せの $i+1$ 桁目が $0$ となれば、その組み合わせは $a+b \lt L$ となることが確定します。よって、 $dp0 _ {i+1} = dp0 _ i \times 3 + dp1 _ i \times 1$ となります。
同様に、 $dp1 _ {i}$ に含まれる組合せの $i+1$ 桁目が $1$ であれば、その組み合わせは $a + b \leqq L$ となる可能性が残ります。よって、 $dp1 _ {i+1} = dp1 _ i \times 2$ となります。

以上を $L$ の最後の桁まで繰り返すことで答えを求めることができます。 $L$ を二進数表記した際の桁数を $N$ とすると、この実装の計算量は $O(N)$ となります。

実装

F - Takahashi's Basics in Education and Learning

使用テクニック： 内積

方針

方針は公式解説通りになります。
ポイントとしては、等差数列を連結した整数を単純な行列の内積で表現することです。等差数列の $i$ 番目までの要素を連結した整数を $X _ i$ 、次に連結される等差数列の要素を $s _ {i+1}$ 、等差数列の公差を $B$ 、連結される等差数列の桁数を $d _ {i+1}$ とすると次の内積を用いた漸化式が成立します。

$\begin{align} \begin{pmatrix}X_i & s_{i+1} & 1\end{pmatrix} \cdot \begin{pmatrix}10^{d_{i+1}} & 0 & 0 \\ 1 & 1& 0 \\ 0 & B & 0\end{pmatrix} &= \begin{pmatrix}X_i \times 10^{d_{i+1}} + s_{i+1} & s_{i+1} + B & 1\end{pmatrix} \\ &= \begin{pmatrix}X_{i+1} & s_{i+2} & 1\end{pmatrix} \end{align}$

問題文の制約から等差数列の要素は $10^{18}$ 未満であるため、桁数が $d$ である等差数列の要素数を $C _ d$ とすると次の式が成立します。

$\begin{pmatrix}X_L & s_{L-1} & 1\end{pmatrix} = \begin{pmatrix}X_1 & s_0 & 1\end{pmatrix} \cdot \begin{pmatrix}10^1 & 0 & 0 \\ 1 & 1& 0 \\ 0 & B & 0\end{pmatrix}^{C_1} \cdots \begin{pmatrix}10^{18} & 0 & 0 \\ 1 & 1& 0 \\ 0 & B & 0\end{pmatrix}^{C_{18}}$