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

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

A - Spoiler

問題

文字列 $S$ が与えられます。
$S$ には $2$ つ | の文字が入っています。
$2$ つの | とその間の文字を $S$ から削除した新たな文字列を出力してください。

考察

for文をつかって文字列 $S$ の文字を左からひとつずつ見ていきます。
このとき、 | の出現回数をカウントしながら進めます。
| に挟まれた文字を見ているとき、 | の出現回数は必ず $1$ 回になっているはずです。
なので、| の出現回数が $1$ 回ではないときの文字を左から順につなぎ合わせればokです。

コード

S = input()

ans = ""
stick_count = 0
for ch in S:
    if ch == '|':
        stick_count += 1
    elif stick_count != 1:
        ans = ans + ch

print(ans)

別解

スライスをつかった解法

まず $2$ つの | のインデックスを取得します。

これはfor文を用いて取得してもいいですし、 index() と rindex() を用いてそれぞれ左側と右側の | のインデックスを取得できるので、そうしてもいいです。

あとはPythonのスライス表記で左サイドと右サイドの削除されない文字列を取得し、 + でその $2$ つをくっつけて出力すればACできます。

S = input()

i, j = S.index('|'), S.rindex('|')
ans = S[:i] + S[j + 1:]
print(ans)

split()メソッドをつかった解法

普段入力の受け取りでお世話になっている .split() ですが、この引数に文字を指定することで、その文字で文字列を分割することができます。

""".split()の使用例"""

S = "AtCoder"
left, right = S.split("o")
print(left)  # AtC
print(right)  # der

これをつかうと、かなり簡潔にコードを書くことができます(これで省略できる問題がA問題でたまに出るので覚えとくといいかもです)。

a, b, c = input().split('|')
print(a + c)

B - Delimiter

問題

$N$ 個の整数 $A_1, A_2, \cdots , A_N$ が $N$ 行にわたって与えられます。
ただし、入力で $N$ は与えられません。
最後の整数 $A_N$ の値は必ず $0$ で、それ以外の入力 $A_1, A_2, \cdots , A_{N-1}$ はどれも $0$ ではありません。
$A_N, A_{N-1}, \cdots , A_1$ の順に $N$ 行にわたって整数を出力してください。

考察

答えとなるリストを用意します。
while文をつかい、 入力を受け取ったらそのリストの末尾に整数をそのまま挿入します。
もしもその整数が $0$ であればそのwhile文を打ち切り、先ほどのリスト内の整数を逆順に出力すればokです。

コード

ans_lst = []
while True:
    a = int(input())
    ans_lst.append(a)
    if a == 0:
        break

for ans in reversed(ans_lst):
    print(ans)

C - A+B+C

問題

$3$ つのリスト $A, B, C$ が与えられます。
$i=1, 2, \cdots , Q$ について、以下の問いに答えてください。

• $3$ つのリスト $A, B, C$ からそれぞれ $1$ つずつ要素を取り出し、その総和を $X_i$ にすることができますか？

考察

制約から、 $3$ つのリスト $A, B, C$ の長さはどれも $100$ 以下です。
よって、$3$ つのリスト $A, B, C$ からそれぞれ $1$ つずつ要素を取り出したときの $3$ つの数の総和は、どんなに多くても $100^3 = 10^6$ 通りしかありません。

総和としてありえる数を事前にすべて列挙してsetに入れておくことで、 $Q$ 個のクエリに $1$ つあたり $O(1)$ で答えられます。

事前のsetの用意が $O(N \times M \times L)$ 、クエリにこたえるパートで $O(Q)$ かかり、全体の計算量は $O(N \times M \times L + Q)$ です。

コード

N = int(input())
A = list(map(int, input().split()))
M = int(input())
B = list(map(int, input().split()))
L = int(input())
C = list(map(int, input().split()))

se = set()
for a in A:
    for b in B:
        for c in C:
            se.add(a + b + c)

Q = int(input())
X = list(map(int, input().split()))
for x in X:
    if x in se:
        print("Yes")
    else:
        print("No")

D - String Bags

問題

空文字列 $S$ があります。
$i=1, 2, \cdots , N$ について、$i$ 番目の袋の中には $A_i$ 個の文字列 $S_{i, 1}, S_{i, 2}, \cdots , S_{i_{A_i}}$ が入っています。
$i=1, 2, \cdots , N$ の順に以下の行動どちらかを選択します。

• $i$ 番目の袋の中の文字列いずれか $1$ を選び、文字列 $S$ の末尾にくっつける。この行動には $1$ 円が必要である
• 何もしない

文字列 $S, T$ を一致させるために必要な最小の金額を求めてください。一致させられない場合は -1 を出力してください。

考察

$$dp[x]: 文字列Tのx文字目までを一致させるのに必要な最小の金額$$
として、$1$ つめの袋から順にこのリストを更新していくとACできます。

細かい説明は省きますが、計算量は全体で $O(N \times \max (A) \times \max (|S|) \times |T|)$ になります。制約より、 $1 \leq N,|T| \leq 100, \enspace 1 \leq \max (A), \max (|S|) \leq 10$ なので、全体で $10^6$ 回程度です。

コード

INF = 1 << 60

T = input()
N = int(input())

dist = [INF] * (len(T) + 1)
dist[0] = 0
for _ in range(N):
    a, *S = input().split()
    nex = dist[:]
    for s in S:
        for i, d in enumerate(dist):
            if d == INF:
                continue
            if s == T[i:i + len(s)]:
                nex[i + len(s)] = min(nex[i + len(s)], dist[i] + 1)
    dist = nex

if dist[-1] == INF:
    print(-1)
else:
    print(dist[-1])

E - Insert or Erase

問題

長さ $N$ の数列 $A=(A_1, A_2, \cdots , A_N)$ があります。
$Q$ このクエリを順に処理してください。クエリは以下の $2$ 種類どちらかです。

• 1 x y : 数列 $A$ 内の $x$ の直後に $y$ を挿入する
• 2 x : 数列 $A$ 内の $x$ を削除する

どの時点でも数列 $A$ 内の要素はどれも互いに異なります。
最終的な数列 $A$ を出力してください。

考察

双方向連結リストとよばれるデータ構造をつかいます。
以下の $2$ つの辞書を管理しながら、各クエリを処理していきます。

• $prev[v]$ : $A$ 内の要素 $v$ の前にある要素
• $nex[v]$ : $A$ 内の要素 $v$ の後にある要素

この方法はリスト内に同じ要素が複数あるとダメですが、今回は常に数列 $A$ 内の要素はどれも互いに異なることが保証されているのでokです。

下のコードでは実装を楽にするために、数列 $A$ の先頭と末尾に $-1, -2$ をそれぞれ追加しています。

コード

N = int(input())
A = [-1] + list(map(int, input().split())) + [-2]

nex, prev = dict(), dict()
for i in range(N + 1):
    nex[A[i]] = A[i + 1]
    prev[A[i + 1]] = A[i]

Q = int(input())
for _ in range(Q):
    q = list(map(int, input().split()))
    if q[0] == 1:
        # x → z(=nex[x]) の並びが、 x → y → z になる。
        x, y = q[1], q[2]
        z = nex[x]
        nex[x], prev[y] = y, x
        nex[y], prev[z] = z, y
    else:
        # y(=prev[x]) → x → z(=nex[x]) の並びが、 y → z になる。
        x = q[1]
        y, z = prev[x], nex[x]
        nex[y], prev[z] = z, y

v = nex[-1]
ans_lst = []
while v != -2:
    ans_lst.append(v)
    v = nex[v]

print(*ans_lst)

おまけ

競プロとは関係ありませんが、基本情報技術者試験というIT国家試験では、この双方向連結リストは良く出題されるみたいです。