AtCoder Beginners Contest 326 (ABC326) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。
A - 2UP3DOWN
問題
高橋くんは短い範囲なら階段をつかって、そうじゃなかったらエレベータをつかいます。
$X$ 階から $Y$ 階に移動するとき、階段をつかいますか?
考察
$a$ 階分のぼったことを $+a$ 、 $b$ 階分おりたことを $-b$ と表すことにします。
問題では、「高橋君は $2$ 階分までの上り、または、$3$ 階分までの下りであれば移動には階段を使い、」 と書いてあるので、 のぼりおりの数値が $-3$ 以上 $2$ 以下なら階段をつかう(=Yesを出力する)ことになります。
だいたいの言語だと -3<=Y-X and Y-X<=2
みたいに式を分けて書かないといけないですが、Pythonだと -3<=Y-X<=2
と1つの式で書けます!(えらい!)
コード
X, Y = map(int, input().split())
if -3 <= Y - X <= 2:
print("Yes")
else:
print("No")
B - 326-like Numbers
問題
$N$ 以上で最小の326-like number(← $3 \times 2 = 6$ みたいになる3桁の数) はいくつですか?
考察
先にこの小問題が解けるかを考えてみます。
$3$ 桁の正整数 $x$ が与えられます。 $x$ は326-like number ですか?
326-like number は、$(100の位の数) \times (10の位の数)=(1の位の数)$ になる数です。
なので、この小問題は、次のようにして解けます。
x = input() # 数値ではなく、文字列で受け取る。
a, b, c = int(x[0]), int(x[1]), int(x[2])
if a * b == c:
print("Yes")
else:
print("No")
これを、$x=N$ のとき、 $x=N+1$ のとき、 $x=N+2$ のとき、 $\cdots$ と繰り返していき、最初の326-like numberを出力すればいいです。
先ほどの小問題のコードを関数化しておきます。
def is_326_like_number(x):
a, b, c = int(x[0]), int(x[1]), int(x[2])
return a * b == c
あとは、for文で $x=N, N+1, N+2, \cdots$ と繰り返していく中で、326-like number が見つかった瞬間にその数を出力しておしまいです。
コード
def is_326_like_number(x):
a, b, c = int(x[0]), int(x[1]), int(x[2])
return a * b == c
N = int(input())
for i in range(N, 1000):
x = str(i) # 文字列に変換してからis_326_like_number(x)をつかう
if is_326_like_number(x):
print(x)
break
C - Peak
問題
長さ $M$ の区間を選びます、その区間にあるプレゼントをすべて受け取れます。
最大で何個のプレゼントをゲットできますか?
考察
リスト $A$ の順番は関係ありません。 たとえば、 $A=(3,5,8)$ と $A=(8,3,5)$ の答えは同じになっているはずです。
なので、いったん数直線に並んでいる順番にリスト $A$ をソートします。
サンプル $1$ の解説では、半開区間 $[1.5,7.5)$ を指定して座標 $2,3,5,7$ にある計 $4$ つのプレゼントを獲得すると書いてありますが、これはあくまで一例です。
たしかに $[1.5,7.5)$ を選択してもいいですが、左側も右側も $0.5$ ずつあいていて、ちょっと効率がわるそうに見えます。
なので、左側を詰めることにします。
つまり、たとえば $A=(2,3,5,7,11,\cdots)$ だと、 $[2, 8)$ の区間、 $[3, 9)$ の区間、 $[5, 11)$ の区間、 $[7, 13)$ の区間、 $[11, 17)$ の区間、 $\cdots$ と順番に見ていくことにします。
最初の $[2,8)$ の区間を選んだときに、プレゼントをいくつ獲得できるかを考えてみます。
区間の右端である $8$ が、リスト $A$ の中でどこに位置するのかが分かるとうれしいです。これが分かると、インデックス同士の引き算で、獲得できるプレゼントの個数が分かります。
「ソートされたリストの中で、この数が何番目に入るのか知りたい」、これは2分探索で解くことができます。
最初のリスト $A$ のソートが $O(NlogN)$ 。
2分探索で獲得できるプレゼントの個数を求めるのが1回あたり $O(logN)$ でこれを $N$ 回するので $O(NlogN)$ 。
全体として $O(NlogN)$ で解くことができました。
コード
from bisect import bisect_left
N, M = map(int, input().split())
A = list(map(int, input().split()))
A.sort()
ans = 0
for left_idx, el in enumerate(A):
right_idx = bisect_left(A, el + M)
ans = max(ans, right_idx - left_idx)
print(ans)
D - ABC Puzzle
問題
いい感じに 'A', 'B', 'C' を配置してね。
https://atcoder.jp/contests/abc326/tasks/abc326_d
考察
問題では 1辺のマスの数 $N$ は $3 \leqq N \leqq 5$ ですが、一番難しそうな $N=5$ で考えてみます。
お気持ちとしては、
- $5 \times 5$ のマスを用意する。
- 各行・各列に1つずつ 'A'があるように、 $5$ つ 'A' を配置する。
- 各行・各列に1つずつ 'B'があるように、 $5$ つ 'B' を配置する。
- 各行・各列に1つずつ 'C'があるように、 $5$ つ 'C' を配置する。
- 各行で、左から順に見たときの最初の文字があっているかどうか確認する。
- 各列で、上から順に見たときの最初の文字があっているかどうか確認する。
の順番で解いていきます。
まず「2. 各行・各列に1つずつ 'A'があるように、 $5$ つ 'A' を配置する。」の部分を考えます。
たとえば、 $1$ 行目は左から $3$ 番目に、 $2$ 行目は左から $1$ 番目に、 $\cdots$ と'A' を配置することを、 $(3,1,4,5,2)$ と数列で表すことにします。
これは、 $P=(1,2,3,4,5)$ の数列を並びかえたものを考えるのと同じです。
itertools.permutations()
をつかうと、順列の並び替えを全パターン探索できます(詳しくは下のコードを見てね)。
'B', 'C' についても同じように配置します。
ただし、配置する文字の位置がかぶっていたらダメなので、そのときは無視します。
「5. 各行で、左から順に見たときの最初の文字があっているかどうか確認する。」の部分は、特に何の工夫もなく実装して大丈夫です。
$5$ 行のうち $3$ 行は左から $1$ 番目に何か文字が書いてあるし、残りの $2$ 行も、最悪でも $3$ 番目には何か文字が書いてあって、かなり軽い探索です。
$4$ 番までの 'A', 'B', 'C' の文字を埋めるところで、 $(5!)^3=1728000 < 2 \times 10^6$ の計算をして、さらにそこから $5, 6$ 番のチェックが入って若干重いですが、実行時間の制限が 4sec になっているので大丈夫です。
コード
from itertools import permutations
def grid_iterator():
for perm_a in permutations(range(N)):
for perm_b in permutations(range(N)):
for perm_c in permutations(range(N)):
is_ok = True
for aj, bj, cj in zip(perm_a, perm_b, perm_c):
# 同じ位置に文字を書くことはできない
if len({aj, bj, cj}) != 3:
is_ok = False
if is_ok:
grid = [['.'] * N for _ in range(N)]
for i, j in enumerate(perm_a):
grid[i][j] = 'A'
for i, j in enumerate(perm_b):
grid[i][j] = 'B'
for i, j in enumerate(perm_c):
grid[i][j] = 'C'
yield grid
def r_check(grid):
for i in range(N):
for j in range(N):
if grid[i][j] == '.':
continue
if grid[i][j] == R[i]:
break
return False
return True
def c_check(grid):
for j in range(N):
for i in range(N):
if grid[i][j] == '.':
continue
if grid[i][j] == C[j]:
break
return False
return True
N = int(input())
R = input()
C = input()
for grid in grid_iterator():
if r_check(grid) and c_check(grid):
print("Yes")
for row in grid:
print(*row, sep='')
exit()
print("No")
E - Revenge of "The Salary of AtCoder Inc."
問題
前に出した目よりも大きい目を出して、たくさん給料もらっちゃおう!
考察
動的計画法で解きます。
$$dp[i]: iの目が出たときの、支給金額の期待値(過去にもらった分は含めない)$$
とします。
答えは $dp[0]$ になります。
急に $dp[0]$ を求めるのは無理ですが、$dp[N]$ ならすぐに求められます。
$N$ の目を出すと $A_N$ 円だけもらえて、その後は何の目が出ても終了してしまうので、 $dp[N] = A_N$ です。
次に $dp[N-1]$ を求めます。
まず $N-1$ の目を出したので、 $A_{N-1}$ 円だけもらいます。
そして、次に出す目が $1$ ~ $N-1$ のとき、そこでこのゲームは終了します。 次に出す目が $N$ のとき、先ほど求めた $dp[N]$ 円がもらえる金額の期待値です。
まとめると、 $N-1$ の目を出したとき、 $A_{N-1} + dp[N]/N$ 円がもらえる金額の期待値になります。
同様にして、$dp[N-2] = A_{N-2} + (dp[N-1] + dp[N])/N$ と求められます。
これを繰り返せば、 $dp[0]$ が求められます。
$A_0=0$ とすると実装が少し簡単になるので、下のコードではそうしています。
コード
MOD = 998244353
N = int(input())
A = [0] + list(map(int, input().split()))
p = pow(N, -1, MOD) # 1/N
# dp[i]: iの目が出たときの、支給金額の期待値(過去にもらった分は含めない)
dp = [0] * (N + 1)
total = 0 # dp[i+1]以降の値の総和
for i in range(N, -1, -1):
dp[i] = (total * p + A[i]) % MOD
total = (total + dp[i]) % MOD
print(dp[0])