PythonでABC341を解いてみたよ。(A～E問題)

AtCoder Beginners Contest 341 (ABC341) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。

A - Print 341

問題

正整数 $N$ が与えられます。
$N+1$ 個の 1 と $N$ 個の 0 からなる、 0 と 1 が交互に並んだ文字列を出力してください。

考察

Pythonでは、文字列の連結や繰り返しを + や * をつかって簡単に書けます。

"""文字列の連結・繰り返し"""
S = "AtCo"
T = "der"

U = S + T
print(U)  # AtCoder

V = T * 3
print(V)  # derderder

答えの文字列は、1文字目の 1 を無視すれば、 01 の文字列を $N$ 回繰り返しています。
なので、答えの文字列は、'1' + '01' * N と書けます。

コード

N = int(input())

ans = '1' + '01' * N
print(ans)

B - Foreign Exchange

問題

$i=1, 2, \cdots ,N$ について、国 $i$ の通貨を $A_i$ 単位もっています。
また、$i=1, 2, \cdots ,N-1$ について、国 $i$ の通貨を $S_i$ 単位だけ支払えば、国 $i+1$ の通貨を $T_i$ 単位だけ獲得できます。
獲得できる国 $N$ の通貨の最大値はいくつですか？

考察

• 可能な限り、国 $1$ の通貨を国 $2$ の通貨に両替する
• 可能な限り、国 $2$ の通貨を国 $3$ の通貨に両替する
• $\cdots$
• 可能な限り、国 $N-1$ の通貨を国 $N$ の通貨に両替する

と両替していけば、国 $N$ の通貨を最大化できます。

国 $i$ の通貨を国 $i+1$ の通貨に両替するとき、 $x=A_i / s$ (小数点以下切り捨て) とすると、国 $i$ の通貨を $s \times x$ 単位だけ支払って、国 $i+1$ の通貨を $t \times x$ 単位だけ獲得します。

これをそのままコードに書き起こすと、下のようになります。

コード

N = int(input())
A = list(map(int, input().split()))

for i in range(N - 1):
    s, t = map(int, input().split())
    x = A[i] // s
    A[i] -= s * x
    A[i + 1] += t * x

print(A[N - 1])

C - Takahashi Gets Lost

問題

$H$ 行 $W$ 列のグリッドがあります。
長さ $N$ の文字列 $T$ にしたがって上下左右の移動をします。この経路上のマスはすべて . が書かれています。
この条件を満たすスタート地点はいくつありますか？

考察

LRUD でマスの移動方向を指定するとき、これをコードの最初に書いておくと実装がラクになることが多いです。

Vec = {
    'D': (1, 0),
    'U': (-1, 0),
    'R': (0, 1),
    'L': (0, -1)
}

各マスについて実際に文字列 $T$ の通りに移動してみて、経路上のマスがすべて . であるようなものの数をカウントすればいいです。

グリッドは $H$ 行 $W$ 列なので、マスの数は $H \times W$ 個あり、各マスについて長さ $N$ の文字列 $T$ による移動シミュレーションをするので、全体で $O(H \times W \times N)$ の計算量でこの問題を解けます。

制約は $1 \leq H, W, N \leq 500$ で、 $H \times W \times N \leq 500^3 = 1.25 \times 10^8$ です。
計算量としてはかなり厳しめですが、実行制限時間が $3$ 秒なのと、上下左右に $1$ マス移動するだけのシミュレーションなので、ギリギリACできます。

コード

Vec = {
    'D': (1, 0),
    'U': (-1, 0),
    'R': (0, 1),
    'L': (0, -1)
}


def is_ok(si, sj):
    if S[si][sj] == '#':
        return False
    i, j = si, sj
    for t in T:
        di, dj = Vec[t]
        i, j = i + di, j + dj
        if not (0 <= i < H and 0 <= j < W):
            return False
        if S[i][j] == '#':
            return False
    return True


H, W, N = map(int, input().split())
T = input()
S = [input() for _ in range(H)]

ans = 0
for i in range(H):
    for j in range(W):
        if is_ok(i, j): ans += 1
print(ans)

おまけ

リストで書き換えて高速化

上のコードを提出すると、おそらく 2900ms前後でACになると思います。

最初に LRUD の辞書を書きましたが、辞書へのアクセスをリストへのアクセスに変えるだけでかなり高速化できます。

実装はほんの少しだけ増えますが、実行時間制限がギリギリのときは辞書ではなくリストで表現した方がいいかもしれません。

Vec = [
    (0, -1),
    (0, 1),
    (-1, 0),
    (1, 0)
]


def is_ok(si, sj):
    if S[si][sj] == '#':
        return 0
    i, j = si, sj
    for t in T:
        di, dj = Vec[t]
        i, j = i + di, j + dj
        if not (0 <= i < H and 0 <= j < W):
            return 0
        if S[i][j] == '#':
            return 0
    return 1


H, W, N = map(int, input().split())
T = list(input())
for i in range(N):
    T[i] = "LRUD".find(T[i])
S = [input() for _ in range(H)]

ans = 0
for i in range(H):
    for j in range(W):
        ans += is_ok(i, j)
print(ans)

D - Only one of two

問題

正整数 $N,M,K$ が与えられます。 $N \neq M$ が保証されています。
$N$ と $M$ のうち どちらか一方のみ で割り切れる数のうち、小さい方から $K$ 番目の数を求めてください。

考察1: 集合で考える

「どちらか一方」と聞かれるとややこしいので、いったんベン図で表してみます。

• 命題A: $N$ で割り切れる
• 命題B: $M$ で割り切れる
  とします。

「 $N$ と $M$ のうち どちらか一方のみ で割り切れる」というのは、ベン図で表すと次の図の黒塗りの部分です。

考察2: 問題を言いかえる

このようなベン図で急に小さい方から $K$ 番目を求めるのはしんどいですが、黒塗りの部分にある要素の数を求める問題は高校の数学Aでよく出題されます。
黒塗りの部分の要素の数は、
$$(Aの要素の数) + (Bの要素の数) - (A \land B の要素の数) \times 2$$
と表せます。

つまり、今回の問題は次のように言い換えられます。

正整数 $x$ に対して、関数 $f(x), g(x), h(x)$ を次のように定めます。
$f(x) := 1以上 x 以下の正整数の中で、Nで割り切れるものの個数$
$g(x) := 1以上 x 以下の正整数の中で、Mで割り切れるものの個数$
$h(x) := 1以上 x 以下の正整数の中で、lcm(N, M)で割り切れるものの個数$
ここで、$lcm(a, b)$ は $a, b$ 最小公倍数を表すとします。
$f(x) + g(x) - 2h(x) \geq K$ を満たす最小の $x$ を求めてください。

考察3: 二分探索

$x$ の値が大きくなればなるほど、ベン図の黒塗りの部分が大きくなります。
つまり大雑把に考えて、$f(x) + g(x) - 2h(x) \geq K$ の式は $x$ の値が小さいと False で、 $x$ の値を大きくしていくとどこかで True になり、それ以降もずっと True のままです。
このような構図の問題では、二分探索がつかえます。

コード

import math


# NとMのうちちょうど一方のみで割り切れるx以下の数はK個以上ですか？
# これがTrueを返す最小のxがans
def judge(x):
    cnt = x // N + x // M - x // lc * 2
    return cnt >= K


N, M, K = map(int, input().split())
lc = math.lcm(N, M)

ok, ng = 1 << 70, 0
while ok - ng > 1:
    mid = (ok + ng) // 2
    if judge(mid):
        ok = mid
    else:
        ng = mid

print(ok)

E - Alternating String

問題

0 と 1 からなる文字列で、文字列中に連続して同じ文字が出てこないものを 良い文字列 とよびます。
0 と 1 からなる文字列 $S$ が与えられます。　$Q$ このクエリを順に処理してください。クエリは以下の $2$ 種類です。

• 1 L R : $S$ の $L$ 文字目から $R$ 文字目までの 0 と 1 を反転させる
• 2 L R : $S$ の $L$ 文字目から $R$ 文字目までの文字列が 良い文字列 であれば Yes を、そうでなければ No を出力する

考察

01列があって、ある区間の01を反転させるクエリを複数回処理するとき、隣り合った要素の差のリスト(=階差数列) を管理すると計算量が落ちることがあります。
(類題: 典型90 問049 Flip Digits 2（★6）)

今回は、隣り合った文字が同じなら $1$ 、異なるなら $0$ となる長さ $N-1$ のリスト $T$ を用意します。
1 L R のクエリが来たとき、すなわち $L$ 番目から $R$ 番目の01を反転させるとき、先ほど用意したリスト $T$ で $L-1$ 番目と $R$ 番目の01を反転させます。

2 L R のクエリが来たとき、リスト $T$ の $L$ 番目から $R-1$ 番目までの要素の総和が $0$ なら Yes で、そうでなければ No です。

区間和を求めるクエリの処理は、セグメントツリーやフェニックツリーをつかうと1回あたり $O(\log N)$ で求められます。

コード

from atcoder.segtree import SegTree

N, Q = map(int, input().split())
S = input()

seg = SegTree(lambda x, y: x + y, 0, N - 1)
for i in range(N - 1):
    if S[i] == S[i + 1]:
        seg.set(i, 1)

for _ in range(Q):
    t, l, r = map(lambda x: int(x) - 1, input().split())
    if t == 0:
        for i in [l - 1, r]:
            if not 0 <= i < N - 1:
                continue
            if seg.get(i) == 1:
                seg.set(i, 0)
            else:
                seg.set(i, 1)
    elif t == 1:
        if seg.prod(l, r) == 0:
            print("Yes")
        else:
            print("No")