AtCoder Beginners Contest 342 (ABC342) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。
A - Yay!
問題
英小文字からなる文字列 $S$ が与えられます。
$S$ はある $1$ 文字を除いて、すべて同じ文字で構成されています。
他のどの文字とも異なる文字は、前から何文字目ですか？
考察
$N$ を文字列 $S$ の長さとします。
$1$ 文字目、 $2$ 文字目、 $\cdots$ 、 $N$ 文字目それぞれについて、その文字が文字列 $S$ 内で $1$ 回しか出てこないかどうかを確認します。
コード
S = input()
length = len(S)
for i in range(length):
cnt = 0
for j in range(length):
if S[i] == S[j]:
cnt += 1
if cnt == 1:
print(i + 1)
別解
.count()をつかった解法
Pythonには、ある要素が文字列の中で出てくる回数を返すメソッド .count() があります。
""".count()の使用例"""
S = "coffeeeeee"
print(S.count('f')) # 2
print(S.count('e')) # 6
これをつかうとより簡潔に書けます。
S = input()
length = len(S)
for i in range(length):
if S.count(S[i]) == 1:
print(i + 1)
collections.Counter()をつかった解法
collections.Counter をつかうと、文字列の中でどの文字が何回現れるかが分かります。
"""collections.Counterの使用例"""
from collections import Counter
S = "coffeeeeee"
C = Counter(S)
print(C) # Counter({'e': 6, 'f': 2, 'c': 1, 'o': 1})
print(C['e']) # 6
print(C['n']) # 0 (存在していない要素もキーとして指定でき、そのときは0を返す)
これをつかうと簡潔に書けます。
計算量的にも、上2つの解法は $O(N^2)$ なのに対し、この解法はCounterの作成に $O(N)$ かかるだけでCounter内の検索は $O(1)$ なので全体として $O(N)$ で少し優秀です。
今回の問題では文字列 $S$ の長さが短く計算量を気にする必要はないですが、より難しい問題を扱うときにつかうことがあるかもしれません。
from collections import Counter
S = input()
C = Counter(S)
for i in range(len(S)):
if C[S[i]] == 1:
print(i + 1)
B - Which is ahead?
問題
$N$ 人が $1$ 列に並んでいます。前から $i$ 番目には人 $P_i$ が立っています。
$Q$ 個のクエリを順に処理してください。 $i$ 番目のクエリは以下の通りです。
- 人 $A_i$ と人 $B_i$ で、前に立っている側の人の番号を出力してください
考察
クエリが来るごとに、人 $A_i, B_i$ それぞれの位置をfor文をつかって求めます。
コード
N = int(input())
P = list(map(int, input().split()))
Q = int(input())
for _ in range(Q):
A, B = map(int, input().split())
ai, bi = -1, -1
for i in range(N):
if P[i] == A:
ai = i
elif P[i] == B:
bi = i
if ai < bi:
print(A)
else:
print(B)
別解
.index()をつかった解法
Pythonには、要素がリスト内ではじめて出現する位置を求めるメソッド .index() があります。
""".index()の使用例"""
array = [3, 1, 4, 1, 5]
print(array.index(4)) # 2
print(array.index(1)) # 1 (最初に出てきたインデックスを返す)
print(array.index(9)) # (ValueErrorを起こす)
これをつかうと、下のように書けます。
N = int(input())
P = list(map(int, input().split()))
Q = int(input())
for _ in range(Q):
A, B = map(int, input().split())
ai, bi = P.index(A), P.index(B)
if ai < bi:
print(A)
else:
print(B)
O(N)解法
$$A_i = 人iが立っている位置$$
となるリスト $A$ を用意してクエリにこたえていくことで、ACできます。
実装を楽にするため、コード内は0-indexedで書いて出力するとき1-indexedに戻しています。
N = int(input())
P = list(map(lambda x: int(x) - 1, input().split()))
# A[i]: 人iの位置
A = [-1] * N
for i in range(N):
A[P[i]] = i
Q = int(input())
for _ in range(Q):
a, b = map(lambda x: int(x) - 1, input().split())
if A[a] < A[b]:
print(a + 1)
else:
print(b + 1)
C - Many Replacement
問題
長さ $N$ の文字列 $S$ が与えられます。
$Q$ 個のクエリを順に処理し、最終的な文字列 $S$ を出力してください。 $i$ 番目のクエリは以下の通りです。
- 文字列 $S$ 内の文字 $c_i$ をすべて $d_i$ に変換してください
考察1 (愚直な実装)
問題文の通りに実装してみると、次のような感じになると思います。
N = int(input())
S = input()
Q = int(input())
for _ in range(Q):
c, d = input().split()
S = S.replace(c, d)
print(S)
計算量を考えてみます。
コードの下から $2$ 行目の S = S.replace(c, d) の部分が $O(N)$ で、この操作を計 $Q$ 回しているので、全体で $O(NQ)$ です。
制約は $1 \leq N, Q \leq 2 \times 10^5$ なので、$O(NQ)$ の計算量だと実行時間制限に間に合いません。
考察2 (辞書の作成)
各クエリでちゃんと文字列 $S$ を変更するのではなく、キーが元々の文字(1文字)・バリューが $Q$ 回の操作後の文字(1文字)となる辞書を用意し、各アルファベットがどのように変更されるのかが分かるようにします。各クエリではこの辞書のみを変更することにします。
これならアルファベットの種類数を $\sigma$ とすると、計算量は全体で $O(N + Q \times \sigma)$ となり、実行時間制限に間に合います。
コード
import string
N = int(input())
S = input()
dic = {ch: ch for ch in string.ascii_lowercase}
Q = int(input())
for _ in range(Q):
c, d = input().split()
for ch in string.ascii_lowercase:
if dic[ch] == c:
dic[ch] = d
ans_lst = [dic[s] for s in S]
print(*ans_lst, sep='')
D - Square Pair
問題
長さ $N$ の非負整数列 $A=(A_1,A_2, \cdots ,A_N)$ について、以下の条件をすべて満たすような整数の組 $(i,j)$ の個数を求めてください。
- $1 \leq i \lt j \leq N$
- $A_i \times A_j$ は平方数( $0$ を含む)
考察1 (平方数の性質)
平方数を列挙すると、$0^2, 1^2, 2^2, 3^2, 4^2, \cdots$ です。
つまり素因数分解をしたとき、たとえば $144 = 2^4 \times 3^2$ について、指数(右上の数)がすべて偶数ならそれは平方数で、そうでなければ平方数ではありません。なので $144$ は平方数です。
考察2 (2つの数の積が平方数になる条件)
$n$ に $1$ 以上の整数 $x$ をかけるとき、かけた後の数が平方数になるような 最小の $x$ を $f(n)$ とします。
たとえば $n = 1350$ のとき、 $1350 = 2 \times 3^3 \times 5^2$ なので $f(n) = 2 \times 3 = 6$ です。
書きかえると、 $1350 = f(1350) \times 3^2 \times 5^2$ です。
この性質として、 $f(a) = f(b)$ のとき、 $a \times b$ は平方数になります。
なぜなら、
- $a = f(a) \times ■^2$
- $b = f(b) \times ▲^2$
のように書けるはずで、
$$a \times b = (f(a)f(b)) \times ■^2 \times ▲^2 = f(a)^2 \times ■^2 \times ▲^2$$ と変形できるからです。
素因数分解のコードは次のように書けます。もし手元に素因数分解のライブラリを持っていない人は、コピペしてどこかのファイルに保存しておくことをオススメします。
# 素因数分解 O(√n) (自作のものです、コピペはご自由に！)
def factorization(n):
now_num = n
prime_cnt_list = []
for i in range(2, n + 1):
if i * i > now_num:
break
if now_num % i:
continue
cnt = 0
while not now_num % i:
now_num //= i
cnt += 1
prime_cnt_list.append((i, cnt))
if now_num != 1:
prime_cnt_list.append((now_num, 1))
return prime_cnt_list
考察3 (問題の言い換え)
考察2で導入した関数 $f$ を用いて、今回の問題は次のように言い換えられます。
長さ $N$ の非負整数列 $A=(A_1,A_2, \cdots ,A_N)$ について、以下の条件をすべて満たすような整数の組 $(i,j)$ の個数を求めてください。
- $1 \leq i \lt j \leq N$
- $f(A_i)=f(A_j)$
ただしコーナーケースとして $A_i=0$ のとき、 $A_i \times A_j$ の値は必ず $0$ になり平方数となるので、正しくは次のように言い換えられます。
長さ $N$ の非負整数列 $A=(A_1,A_2, \cdots ,A_N)$ について、以下の2条件をすべて満たすような整数の組 $(i,j)$ の個数を求めてください。
- $1 \leq i \lt j \leq N$
- $f(A_i)=f(A_j)$ または $A_i=0$ または $A_j=0$
考察4 (解く)
$j=1, 2, \cdots N$ と順番に見ていきます。
その中で $1 \leq i \lt j$ となる整数 $i$ すべてについて $f(A_i)$ の値の個数を格納した辞書を用意し管理しながら進めることで、この問題を解くことができます。
(詳しくは下のコードを見てください。)
コード
from collections import defaultdict
# 素因数分解 O(√n)
def factorization(n):
now_num = n
prime_cnt_list = []
for i in range(2, n + 1):
if i * i > now_num:
break
if now_num % i:
continue
cnt = 0
while not now_num % i:
now_num //= i
cnt += 1
prime_cnt_list.append((i, cnt))
if now_num != 1:
prime_cnt_list.append((now_num, 1))
return prime_cnt_list
def f(x):
fc = factorization(x)
result = 1
for prime, cnt in fc:
if cnt % 2 == 1:
result *= prime
return result
N = int(input())
A = list(map(int, input().split()))
dic = defaultdict(int) # i<j となるすべてのiについて、A_iの個数を管理する辞書
ans = 0
for j, a in enumerate(A):
if a == 0:
dic[0] += 1
ans += j
else:
val = f(a)
ans += dic[val] + dic[0]
dic[val] += 1
print(ans)
別解
線形篩をつかった解法
素因数分解をするパートについて、
$$P_x = (xに含まれる2以上の最小の素因数)$$
として、これを $x = 2$ から $x=2 \times 10^5$ まで用意します。構築はエラトステネスの篩と同じ要領で $O(N \log \log N)$ でできて、素因数分解は $O(\log N)$ でできます。
「素因数分解を複数回する + 素因数分解する値の最大値がそこまで大きくない」 の 2条件を満たすときにつかえて、普通の素因数分解よりも高速です。
from collections import defaultdict
class LinearSieve:
def __init__(self, n_max):
self.n_max = n_max
self.__build()
def __build(self):
self.arr = [i for i in range(self.n_max + 1)]
self.arr[0] = -1
self.arr[1] = -1
for i in range(2, self.n_max + 1):
# これ以降は、iが最小の素因数になることはないので、breakする。
if i * i > self.n_max:
break
# iが素数でないときはcontinue
if self.arr[i] != i:
continue
num = i + i
while num <= self.n_max:
if self.arr[num] == num:
self.arr[num] = i
num += i
def is_prime(self, num: int) -> bool:
assert 1 <= num <= self.n_max
return self.arr[num] == num
def factorization(self, num: int) -> list:
"""素因数分解
O(log(num))
Args:
num: 素因数分解したい数
Returns: [素因数, その個数] のリスト
"""
assert 1 <= num <= self.n_max
min_divisor = -2
ret = []
while num > 1:
divisor = self.arr[num]
if min_divisor == divisor:
ret[-1][1] += 1
else:
ret.append([divisor, 1])
min_divisor = divisor
num //= divisor
return ret
N = int(input())
A = list(map(int, input().split()))
dic = defaultdict(int)
ls = LinearSieve(200_003)
ans = 0
for i, a in enumerate(A):
if a == 0:
ans += i
dic[0] += 1
else:
fc = ls.factorization(a)
v = 1
for divisor, cnt in fc:
if cnt & 1:
v *= divisor
ans += dic[v]
ans += dic[0]
dic[v] += 1
print(ans)
E - Last Train
問題
駅 $1$ 、駅 $2$ 、 $\cdots$ 、駅 $N$ の計 $N$ 個の駅があります。
また電車についての情報は $M$ 個あり、 $i$ 個目の情報は以下の通りです。
- $t = l_i, \enspace l_i + d_i, \enspace l_i + 2d_i, \enspace \cdots , \enspace l_i + (k_i - 1)d_i$ について、次のような電車が存在する
- 時刻 $t$ に駅 $A_i$ を出発し、時刻 $t + c_i$ に駅 $B_i$ に到着する
ゴール地点を駅 $N$ とするとき、 $i = 1, 2, \cdots N-1$ について、駅 $i$ の終電はいつですか？
考察
ダイクストラ法で解きます。
$$dp[i] = (駅 i の終電の時刻)$$
とします。駅 $N$ の終電の時刻を便宜上 $\infty$ とし、駅 $N$ からこのリスト $dp$ を完成させていきます。
そうすると電車の情報 $(l_i,d,_ik_i,c_i,A_i,B_i)$ について、到着地点の駅 $B_i$ の終電の時刻が分かったときに出発地点の駅 $A_i$ の終電の時刻を更新したくなります。なので、入力は次のように受け取ります。
N, M = map(int, input().split())
G_rev = [[] for _ in range(N)]
for _ in range(M):
l, d, k, c, A, B = map(int, input().split())
A, B = A - 1, B - 1
G_rev[B].append((l, d, k, c, A))
駅 $B_i$ の終電の時刻が $t_{B_i}$ のとき、駅 $A_i$ の終電の時刻 $t_{A_i}$ を計算する方法について考えます。
$x=0,1, \cdots N-1$ について、電車の発車する時刻は $t_1=l_i + x \times d_i$ で、到着時刻は $t_2 = l_i + x \times d_i + c_i$ です。
駅 $B_i$ の終電が $t_{B_i}$ だと分かっているので、次の式を満たしている必要があります。
$$t_2 = l_i + x \times d_i + c_i \leq t_{B_i}$$
この式を満たす最大の $x$ を求めれば、乗るべき電車が分かります。
$x$ について整理すると、$x = (t_{B_i} - l_i - c_i) / d_i$ (小数点以下切り捨て) です。
これで $x$ が分かったので乗るべき電車が分かりました。あとは優先度付きキューを利用してダイクストラ法をつかってACできます。
コード
from heapq import heappop, heappush
INF = 1 << 63
N, M = map(int, input().split())
G_rev = [[] for _ in range(N)]
for _ in range(M):
l, d, k, c, A, B = map(int, input().split())
A, B = A - 1, B - 1
G_rev[B].append((l, d, k, c, A))
# dp[station]: last_time
dp = [-INF] * N
dp[-1] = INF
que = []
heappush(que, (-INF, N - 1)) # (-last_time, station)
while que:
mt, station = heappop(que)
pt = -mt
if dp[station] != pt:
continue
for l, d, k, c, A in G_rev[station]:
x = min((pt - l - c) // d, k - 1)
if x < 0:
continue
nt = l + x * d
if dp[A] < nt:
dp[A] = nt
heappush(que, (-nt, A))
for i in range(N - 1):
ans = dp[i]
if ans == -INF:
ans = "Unreachable"
print(ans)