AtCoder Beginners Contest 334 (ABC334) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。
A - Christmas Present
問題
サンタさんはバットかグローブのうち値段が高い方を高橋くんにプレゼントします。
バットとグローブの値段がそれぞれ $B$ 円、 $G$ 円( $B\neq G$ )のとき、プレゼントするのはどちらですか？
考察
if文をつかって $B < G$ のとき "Bat" と出力し、そうでないときは "Glove" と出力します。
コード
B, G = map(int, input().split())
if B > G:
print("Bat")
else:
print("Glove")
B - Christmas Trees
問題
ある整数 $k$ を用いて $A+kM$ と表されるような地点それぞれにクリスマスツリーがあります。( $A, M$ は入力で与えられる整数で、 $k$ は任意の整数です )
高橋くんと青木くんがそれぞれ 座標 $L, R$ にいるとき、この $2$ 人の間にあるクリスマスツリーはいくつありますか？
考察
図を用いながら見ていきます。
クリスマスツリーはこんな感じで並んでいます。
上の図では、 $A+kM$ の $k$ の値をそのままクリスマスツリーの番号とすることにしています。
高橋くんと青木くんの間にあるクリスマスツリーの個数を求めたいので、座標 $L$ にいる高橋くんのすぐ右側にあるクリスマスツリーの番号と、座標 $R$ にいる青木くんのすぐ左側にあるクリスマスツリーの番号を調べることにします。
たとえば、高橋くんのすぐ右側にあるクリスマスツリーの番号が $3$ で、青木くんのすぐ左側にあるクリスマスツリーの番号が $10$ だったら、 $3～10$ のクリスマスツリーの個数である $8$ が答えになります。
さっそく、高橋くんのすぐ右側にあるクリスマスツリーの番号を調べます。
たとえばこの図だと、高橋くんのすぐ右側にあるクリスマスツリーの番号は $1$ です。
式をつかって表すと、 $L \leq A+kM$ が成り立つような $k$ の最小値が $1$ ということになります。
これを $k$ について解くと、
$$k \geq \frac{L-A}{M}$$
となります。
たとえば、右辺の $\frac{L-A}{M}$ の値がちょうど $2$ だったら、 $k$ の最小値は $2$ ですが、右辺の値が $2.01$ になると、 $k$ の最小値は $3$ になります。
つまり、小数点以下を切り上げるような割り算をすることになります。
コードで書くとこんな感じ。
"""小数点以下切り捨ての割り算をした後、もし小数点以下があるなら+1する"""
left = (L - A) // M
if (L - A) % M != 0:
left += 1
青木くんのすぐ左側にあるクリスマスツリーの番号についても同じように考えていきます。
上の図を式で表すと、 $A+kM \leq R$ が成り立つような $k$ の最大値が $1$ だということになります。
式変形するとこう。
$$k \leq \frac{R-A}{M}$$
詳細は省きますが、先ほどとちがって今回は小数点以下切り捨ての割り算をすればいいので、さっきみたいな工夫は必要ありません。
これで、高橋くんと青木くんの間にあるクリスマスツリーについて、左端と右端のクリスマスツリーの番号が分かりました。よって、このクリスマスツリーの個数が求められました。
コード
A, M, L, R = map(int, input().split())
left = (L - A) // M
if (L - A) % M != 0:
left += 1
right = (R - A) // M
print(right - left + 1)
おまけ
小数点以下切り上げの割り算について
次のコードの方が簡潔で計算回数も少なく済みます。この記事では可読性も重視しているので、ここで紹介することにしました。
# a÷bの小数点以下切り上げ
def ceil_div(a, b):
return (a + b - 1) // b
C - Socks 2
問題
$i=1,2,\cdots N$ について、色 $i$ の靴下が $2$ 枚あります。
ただし、色 $A_1, A_2, \cdots ,A_K$ の靴下は、かたっぽだけ紛失して $1$ 枚だけになっています。
この $2N-K$ 枚の靴下をうまく組み合わせて $\lfloor\frac{2N-K}{2}\rfloor$ 個のペアをつくります。色 $i,j$ の靴下を組み合わせたときの「奇妙度」を $|i-j|$ とするとき、奇妙度の総和の最小値を求めてください。
考察1 (靴下の組み合わせ方)
たとえば色 $1, 3, 5$ の靴下がそれぞれ $1, 2, 1$ 枚あったとして、 $(1, 3), (3, 5)$ と組み合わせても $(1, 5), (3, 3)$ と組み合わせても、奇妙度はどちらも $4$ で変わりません。
図にするとこんな感じです。
ということで、揃っている靴下はすべて揃っている色同士で組み合わせてしまいます。
ここから先は、揃っていない靴下のみを考えていきます。
考察2 (Kが偶数のとき)
たとえば色 $(1, 2, 3, 4)$ の靴下をうまく組み合わせたいとき、 $(1, 2), (3, 4)$ と組み合わせるのが最適で、 $(1, 3), (2, 4)$ や $(1, 4), (2, 3)$ と組み合わせると損してしまいます。
$K$ が奇数のときは後でするとして、いったんここまでのコードはこうなります。
N, K = map(int, input().split())
A = list(map(int, input().split()))
if K % 2 == 0:
ans = 0
for i in range(0, K, 2):
ans += A[i + 1] - A[i]
print(ans)
else:
pass # 考察3でやっていきます。
考察3 (Kが奇数のとき)
サンプル $3$ の入力が $A=(1, 2, 4, 7,8)$ で、出力は $2$ です。
これは $(1, 2), (7, 8)$ を組み合わせて、色 $4$ の靴下を無視することで達成できます。
どの靴下を無視するかに注目する代わりに、$i=0,1,\cdots ,K/2$ それぞれについて「左から順番に $i$ 個のペアをつくって、右からも順番に $K/2-i$ 個のペアをつくったときの奇妙度の総和」を考えていくことにします。
左から順にペアをつくったとき、 $i$ 個目のペアの奇妙度が $P_i$ になるとします。
コードで書くと、こうです。
P = []
for i in range(K // 2):
P.append(A[i * 2 + 1] - A[i * 2])
同様に、右から順にペアをつくったとき、 $j$ 個目のペアの奇妙度が $Q_j$ になるとします。
Q = []
for j in reversed(range(K // 2)):
Q.append(A[j * 2 + 2] - A[j * 2 + 1])
左から $i$ 個のペアをつくったときの $i$ 個のペアの奇妙度の総和がほしいので、リスト $P, Q$ それぞれの累積和がほしいです。
累積和は itertools.accumulate() をつかうとラクに実装できます。
from itertools import accumulate
rui_P = [0] + list(accumulate(P))
rui_Q = [0] + list(accumulate(Q))
あとは、$i=0,1,\cdots ,K/2$ それぞれについて「左から順番に $i$ 個のペアをつくって、右からも順番に $K/2-i$ 個のペアをつくったときの奇妙度の総和」を求め、その最小値が答えになります。
ans = 10 ** 17
for i in range(K // 2 + 1):
ans = min(ans, rui_P[i] + rui_Q[K // 2 - i])
print(ans)
まとめると、下のコードになります。
コード
from itertools import accumulate
N, K = map(int, input().split())
A = list(map(int, input().split()))
if K % 2 == 0:
ans = 0
for i in range(0, K, 2):
ans += A[i + 1] - A[i]
print(ans)
else:
P = []
for i in range(K // 2):
P.append(A[i * 2 + 1] - A[i * 2])
Q = []
for j in reversed(range(K // 2)):
Q.append(A[j * 2 + 2] - A[j * 2 + 1])
rui_P = [0] + list(accumulate(P))
rui_Q = [0] + list(accumulate(Q))
ans = 10 ** 17
for i in range(K // 2 + 1):
ans = min(ans, rui_P[i] + rui_Q[K // 2 - i])
print(ans)
D - Reindeer and Sleigh
問題
ソリ $i$ をひくために必要なトナカイは $R_i$ 匹です。
以下の形式のクエリが計 $Q$ 個やってきます。処理してください。
- 整数 $X$ が与えられます、トナカイが $X$ 匹いたとき最大で何匹のソリをひけますか？
考察
「ひくのに必要になるトナカイの数が少ないソリ」から順番にひいていきます。
なので、リスト $R$ を小さい順にソートします。
$k$ 台のソリをひくとき、必要なトナカイは $R_1+R_2+\cdots +R_k$ 匹です。
累積和をつかって、事前に計算しておきます。累積和は itertools.accumulate() をつかうとラクに実装できます。
from itertools import accumulate
rui = list(accumulate(R))
各クエリで rui[i] <= X を満たす最大の $i$ を求めればよいです。
rui[0] <= X を満たすよりも、 rui[N-1] <= X を満たす方がしんどいです。リストの後ろに行くにつれて徐々に条件を満たしにくくなるこの構図は、二分探索がつかえます。
二分探索をつかうとクエリ $1$ 回あたり $O(logN)$ で答えを求められます。
コード
from itertools import accumulate
from bisect import bisect_right
N, Q = map(int, input().split())
R = list(map(int, input().split()))
R.sort()
rui = list(accumulate(R))
for _ in range(Q):
x = int(input())
print(bisect_right(rui, x))
E - Christmas Color Grid 1
問題
$H \times W$ のグリッドがあり、各マスは赤色or緑色で塗られています。
上下左右に隣り合った緑色のマス同士はつながっているとみなします。
赤色のマスのうち $1$ つがランダムに選ばれて緑色に塗られます。このとき緑色のマスの連結成分の個数の期待値はいくつですか？ $\text{mod} 998244353$ で求めてください。
考察
$H \times W \leq 10^6$ なので、赤色のマスを $1$ つ緑色に変えたところで大して連結成分の個数は変わりません。
先にUnion-Findで緑色のマスの連結をしておきます。
ac-library-pythonのUnion-Findは1次元にしか対応していないので、マス $(i, j)$ は頂点 $i*W+j$ とおきかえてコードを書くことにします。
また、各マスを1つずつ見ていくので、ついでに赤色のマスがどこにあるのかも記録しておきます。
from atcoder.dsu import DSU
def pos(i, j):
return i * W + j
red = []
uf = DSU(H * W)
for i in range(H):
for j in range(W):
if S[i][j] == '#':
for ni, nj in [(i + 1, j), (i, j + 1)]:
if not (0 <= ni < H and 0 <= nj < W):
continue
if S[ni][nj] == '#':
uf.merge(pos(i, j), pos(ni, nj))
else:
red.append((i, j))
これで、マスの色を変更する前の連結成分の数を求められます。
prev_cnt = len(uf.groups()) - len(red) # 元々の連結成分の数
ここからは具体的な図を見ながら進めます。
Union-Findの処理をしたので、各連結成分に番号を振って見やすくしてあります。
星マークの付いた赤色のマスを緑色に変えることを考えます。
星マークのマスの上下左右にある数字(連結成分の番号)は $1$ のみです。
ここを緑色に変えても、連結成分は増えません。
その右側のマスも同じです。
星マークの位置を変えました。このマスを緑色に変えることを考えます。
星マークのマスの上下左右にある数字は $1, 2$ の $2$ つです。
このマスを緑色に変えると、番号 $1, 2$ の連結成分がつながり、全体の連結成分の数は $1$ つだけ減ります。
最後です、また星マークの位置を変えました。
このマスの上下左右には $1, 2, 3$ の $3$ つの番号があり、このマスを緑色に変えると $1, 2, 3$ 番の連結成分がつながるため、連結成分の数は $2$ つ減ります。
あとはこれを実装すればokです。
下のコードでは、連結成分の番号を uf.leader()(属している連結成分の代表頂点番号)で代用しています。
コード
from atcoder.dsu import DSU
MOD = 998244353
def pos(i, j):
return i * W + j
H, W = map(int, input().split())
S = [input() for _ in range(H)]
red = []
uf = DSU(H * W)
for i in range(H):
for j in range(W):
if S[i][j] == '#':
for ni, nj in [(i + 1, j), (i, j + 1)]:
if not (0 <= ni < H and 0 <= nj < W):
continue
if S[ni][nj] == '#':
uf.merge(pos(i, j), pos(ni, nj))
else:
red.append((i, j))
prev_cnt = len(uf.groups()) - len(red)
ans = 0
for i, j in red:
se = set()
for ni, nj in [(i + 1, j), (i - 1, j), (i, j + 1), (i, j - 1)]:
if not (0 <= ni < H and 0 <= nj < W):
continue
if S[ni][nj] == '#':
se.add(uf.leader(pos(ni, nj)))
ans += prev_cnt - len(se) + 1
ans %= MOD
ans = ans * pow(len(red), -1, MOD) % MOD # ans/=len(red) をしている
print(ans)