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])