ABC352のA~E問題をpythonで解説していきます。筆者はA-Eの5完でした。
A - AtCoder Line
問題
鉄道の AtCoder 線には $N$ 個の駅があり、それぞれ $1, 2, \ldots, N$ の番号が付けられています。
AtCoder 線では、駅 $1$ を始発駅として駅 $2, 3, \ldots, N$ の順に各駅に停車する 上り列車 および、駅 $N$ を始発駅として駅 $N - 1, N - 2, \ldots, 1$ の順に各駅に停車する 下り列車 が運行されています。
高橋君は AtCoder 線の上り列車あるいは下り列車の一方のみを使うことで駅 $X$ から駅 $Y$ まで移動しようとしています。
この移動の間に高橋君が乗っている電車が駅 $Z$ に停車することがあるか判定してください。
考察
2つの駅について、番号が小さいほうを$X$、大きいほうを$Y$とすれば、後は$Z$がその間にあるかを確かめるだけでよいです。
コード
N, X, Y, Z = map(int, input().split())
if X > Y:
X, Y = Y, X
if X <= Z <= Y:
print("Yes")
else:
print("No")
B - Typing
問題
高橋君は英小文字からなる文字列 $S$ をキーボードで入力しようとしました。
高橋君は画面を見ずにキーボードだけを見てタイピングをしていました。
誤って別の英小文字を入力してしまったときにはその直後にバックスペースキーを押しましたが、バックスペースキーが壊れていたため誤って入力された文字は消去されず、実際に入力された文字列は文字列 $T$ となりました。
また、英小文字以外のキーを誤って押してしまうことはありませんでした。
$T$ のうち高橋君が誤って入力した文字でないものを 正しく入力された文字 であると定めます。
正しく入力された文字が $T$ の何文字目であるか答えてください。
考察
前から一文字ずつ見ていき、$S_{i}$と$T_{j}$が一致したインデックス($j$)を保存して最後にそれを出力すればよいです。
コード
S = list(input())
T = list(input())
i = 0
j = 0
ans = []
while i < len(S) and j < len(T):
if S[i] == T[j]:
i += 1
ans.append(j + 1)
j += 1
print(*ans)
C - Standing On The Shoulders
問題
$N$ 人の巨人がいます。巨人にはそれぞれ $1, 2, \ldots, N$ の名前がついており、巨人 $i$ が地面に立ったとき、肩の高さは $A_i$、頭の高さは $B_i$ となります。
あなたは $(1, 2, \ldots, N)$ を並べ替えて得られる数列 $(P_1, P_2, \ldots, P_N)$ を選び、以下の規則に従って $N$ 人の巨人を積み上げることができます。
- まず地面に巨人 $P_1$ を立たせる。巨人 $P_1$ の肩は地面を基準として $A_{P_1}$、頭は地面を基準として $B_{P_1}$ の高さとなる。
- $i = 1, 2, \ldots, N - 1$ の順に巨人 $P_i$ の肩の上に巨人 $P_{i + 1}$ を立たせる。巨人 $P_i$ の肩が地面を基準として高さ $t$ のとき、巨人 $P_{i + 1}$ の肩は地面を基準として $t + A_{P_{i + 1}}$、頭は地面を基準として $t + B_{P_{i + 1}}$ の高さとなる。
一番上に立っている巨人、すなわち巨人 $P_N$ の地面を基準とした頭の高さとして実現できる最大値を求めてください。
考察
簡単に言ってしまえば、肩車(？)の一番上にいる巨人の目線の高さの最大値を答える問題です。巨人$i$が一番上となるとき、すべての巨人の肩の高さの和を$all$とすると、$all-A_i+B_i$が巨人$i$の目線の高さになります。したがって、すべての$i$について巨人の目線を調べ、その最大値を出力すればよいです。
コード
N = int(input())
A = []
B = []
for _ in range(N):
a, b = map(int, input().split())
A.append(a)
B.append(b)
all = sum(A)
ans = 0
for i in range(N):
ans = max(ans, all - A[i] + B[i])
print(ans)
D - Permutation Subsequence
問題
$(1,2,\dots,N)$を並び替えて得られる数列 $P=(P_1,P_2,\dots,P_N)$ が与えられます。
長さ $K$ の正整数列 $(i_1,i_2,\dots,i_K)$ であって、以下の条件を共に満たすものを良い添字列と呼びます。
- $1\leq i_1 < i_2 < \dots < i_K \leq N$
- $(P_{i_1},P_{i_2},\dots,P_{i_K})$ はある連続する $K$ 個の整数を並び替えることで得られる。
厳密には、ある整数 $a$ が存在して、$\lbrace P_{i_1},P_{i_2},\dots,P_{i_K} \rbrace = \lbrace a,a+1,\dots,a+K-1 \rbrace$。
全ての良い添字列における $i_K-i_1$ の最小値を求めてください。
なお、本問題の制約下では良い添字列が必ず $1$ つ以上存在することが示せます。
考察
$i$から$i+K$まで連続した値がすべて存在するような最小の区間の長さを求める問題です。
このような問題の特徴として、次の$i+1$から$i+K+1$の範囲を考えるとき、前回から$i$の情報を落として$i+K+1$の情報を追加するだけでいいという点が重要になります。(尺取り法に近い)
あとは、情報=その値の存在するインデックスとし、その中の最大値と最小値を高速で求められるようなデータ構造で管理すればよいです。今回はSortedSetを用いて実装しました。
コード
import math
from bisect import bisect_left, bisect_right
from typing import Generic, Iterable, Iterator, TypeVar, Union, List
from collections import defaultdict
T = TypeVar("T")
# https://github.com/tatyam-prime/SortedSet/blob/main/SortedSet.py
class SortedSet(Generic[T]):
BUCKET_RATIO = 50
REBUILD_RATIO = 170
def _build(self, a=None) -> None:
"Evenly divide `a` into buckets."
if a is None:
a = list(self)
size = self.size = len(a)
bucket_size = int(math.ceil(math.sqrt(size / self.BUCKET_RATIO)))
self.a = [
a[size * i // bucket_size : size * (i + 1) // bucket_size]
for i in range(bucket_size)
]
def __init__(self, a: Iterable[T] = []) -> None:
"Make a new SortedSet from iterable. / O(N) if sorted and unique / O(N log N)"
a = list(a)
if not all(a[i] < a[i + 1] for i in range(len(a) - 1)):
a = sorted(set(a))
self._build(a)
def __iter__(self) -> Iterator[T]:
for i in self.a:
for j in i:
yield j
def __reversed__(self) -> Iterator[T]:
for i in reversed(self.a):
for j in reversed(i):
yield j
def __len__(self) -> int:
return self.size
def __repr__(self) -> str:
return "SortedSet" + str(self.a)
def __str__(self) -> str:
s = str(list(self))
return "{" + s[1 : len(s) - 1] + "}"
def _find_bucket(self, x: T) -> List[T]:
"Find the bucket which should contain x. self must not be empty."
for a in self.a:
if x <= a[-1]:
return a
return a
def __contains__(self, x: T) -> bool:
if self.size == 0:
return False
a = self._find_bucket(x)
i = bisect_left(a, x)
return i != len(a) and a[i] == x
def add(self, x: T) -> bool:
"Add an element and return True if added. / O(√N)"
if self.size == 0:
self.a = [[x]]
self.size = 1
return True
a = self._find_bucket(x)
i = bisect_left(a, x)
if i != len(a) and a[i] == x:
return False
a.insert(i, x)
self.size += 1
if len(a) > len(self.a) * self.REBUILD_RATIO:
self._build()
return True
def discard(self, x: T) -> bool:
"Remove an element and return True if removed. / O(√N)"
if self.size == 0:
return False
a = self._find_bucket(x)
i = bisect_left(a, x)
if i == len(a) or a[i] != x:
return False
a.pop(i)
self.size -= 1
if len(a) == 0:
self._build()
return True
def lt(self, x: T) -> Union[T, None]:
"Find the largest element < x, or None if it doesn't exist."
for a in reversed(self.a):
if a[0] < x:
return a[bisect_left(a, x) - 1]
def le(self, x: T) -> Union[T, None]:
"Find the largest element <= x, or None if it doesn't exist."
for a in reversed(self.a):
if a[0] <= x:
return a[bisect_right(a, x) - 1]
def gt(self, x: T) -> Union[T, None]:
"Find the smallest element > x, or None if it doesn't exist."
for a in self.a:
if a[-1] > x:
return a[bisect_right(a, x)]
def ge(self, x: T) -> Union[T, None]:
"Find the smallest element >= x, or None if it doesn't exist."
for a in self.a:
if a[-1] >= x:
return a[bisect_left(a, x)]
def __getitem__(self, x: int) -> T:
"Return the x-th element, or IndexError if it doesn't exist."
if x < 0:
x += self.size
if x < 0:
raise IndexError
for a in self.a:
if x < len(a):
return a[x]
x -= len(a)
raise IndexError
def index(self, x: T) -> int:
"Count the number of elements < x."
ans = 0
for a in self.a:
if a[-1] >= x:
return ans + bisect_left(a, x)
ans += len(a)
return ans
def index_right(self, x: T) -> int:
"Count the number of elements <= x."
ans = 0
for a in self.a:
if a[-1] > x:
return ans + bisect_right(a, x)
ans += len(a)
return ans
INF = float("inf")
N, K = map(int, input().split())
P = list(map(int, input().split()))
s = SortedSet()
d = defaultdict(int)
for i, p in enumerate(P):
d[p] = i
ans = INF
for i in range(1, K + 1):
s.add(d[i])
ans = min(ans, s[-1] - s[0])
for i in range(K + 1, N + 1):
s.discard(d[i - K])
s.add(d[i])
ans = min(ans, s[-1] - s[0])
print(ans)
E - Clique Connect
問題
$N$ 頂点からなる重み付き無向グラフ $G$ があります。 $G$ の各頂点には $1$ から $N$ までの番号が付けられています。 最初、$G$ には辺が $1$ 本も存在しません。
今から、$M$ 回の操作を行うことによって $G$ に辺を追加していきます。 $i\ (1\leq i\leq M)$ 回目の操作は以下の通りです。
- $K_i$ 頂点からなる頂点の部分集合 $S_i=\lbrace A_{i,1},A_{i,2},\dots,A_{i,K_i}\rbrace$ が与えられる。 $u,v\in S_i$ かつ $u<v$ を満たす全ての $u,v$ について、頂点 $u$ と頂点 $v$ の間に重み $C_i$ の辺を追加する。
$M$ 回の操作を全て行ったとき $G$ が連結になるか判定し、連結になるならば $G$ の最小全域木に含まれる辺の重みの総和を求めてください。
考察
最小全域木を考える問題なので、クラスカル法(小さいコストの辺から見ていく)を考え、優先度付きキュー(最小値をpop可能)＋UnionFind(集合演算可能)を使うことを考えます。
ここで、すべての辺をつなぐか考慮するということで、例えば頂点$N$個の場合$O(N^2)$かかりそうと考えるかもしれませんが、これは集合演算なので、ノード1とノード2をつなぐか考え、ノード2とノード3をつなぐか考えた場合、実質的にノード1とノード3をつなぐかは考慮したと考えられる点に注意が必要です。
頂点数の和$\sum_{i=1}^M K_i$は$4*10^5$以下なので、あとは実装するだけでよいです。
コード
import heapq
class UnionFind:
def __init__(self, n):
self.n = n
self.parents = [-1] * n
def find(self, x):
if self.parents[x] < 0:
return x
else:
self.parents[x] = self.find(self.parents[x])
return self.parents[x]
def union(self, x, y):
x = self.find(x)
y = self.find(y)
if x == y:
return
if self.parents[x] > self.parents[y]:
x, y = y, x
self.parents[x] += self.parents[y]
self.parents[y] = x
def size(self, x):
return -self.parents[self.find(x)]
def same(self, x, y):
return self.find(x) == self.find(y)
def members(self, x): # 多用すると重い
root = self.find(x)
return [i for i in range(self.n) if self.find(i) == root]
def roots(self):
return [i for i, x in enumerate(self.parents) if x < 0]
def group_count(self):
return len(self.roots())
def all_group_members(self):
group_members = defaultdict(list)
for member in range(self.n):
group_members[self.find(member)].append(member)
return group_members
def __str__(self):
return "\n".join(f"{r}: {m}" for r, m in self.all_group_members().items())
N, M = map(int, input().split())
pq = []
for _ in range(M):
K, C = map(int, input().split())
A = list(map(int, input().split()))
heapq.heappush(pq, (C, A))
uf = UnionFind(N + 1)
ans = 0
while pq:
cost, A = heapq.heappop(pq)
for i in range(len(A) - 1):
a1, a2 = A[i], A[i + 1]
if uf.same(a1, a2):
continue
uf.union(a1, a2)
ans += cost
if uf.group_count() != 2:
print(-1)
exit()
print(ans)