ABC353のA~E問題をpythonで解説していきます。筆者は予定があって参加できなかったのですが、後日自力で解けたのはA-Eまででした。
A - Buildings
問題
$N$ 個のビルが横一列に並んでいて、左から $i$ 番目のビルの高さは $H_i$ です。
左から $1$ 番目のビルより高いビルが存在するか判定し、存在する場合その内最も左のビルは左から何番目か求めてください。
考察
1番左のビルの高さを基準とし、ビルを左から走査していって基準より高いものが見つかればその1-indexを、なければ-1を返すようにすればよいです。
コード
N = int(input())
H = list(map(int, input().split()))
base = H[0]
for i, h in enumerate(H[1:]):
if base < h:
print(i + 2)
exit()
print(-1)
B - AtCoder Amusement Park
問題
AtCoder 遊園地には $K$ 人乗りのアトラクションがあります。現在、このアトラクションの待機列には $N$ グループが並んでいます。
先頭から $i$ 番目 ($1 \leq i \leq N$) のグループは $A_i$ 人組です。すべての $i$ ($1 \leq i \leq N$) について、$A_i \leq K$ です。
高橋君はこのアトラクションのスタッフとして、並んでいるグループを次の手順に従って誘導します。
はじめ、アトラクションには誰も誘導されておらず、空席は $K$ 個あります。
- 待機列に並んでいるグループがない場合、アトラクションをスタートさせ、誘導を終了する。
- アトラクションの空席の数と待機列の先頭に並んでいるグループの人数を比較し、次のどちらかを行う。
待機列の先頭に並んでいるグループの人数よりアトラクションの空席の数のほうが少ない場合、アトラクションをスタートさせる。スタートしたのち、アトラクションの空席が $K$ 個になる。
そうでない場合、待機列の先頭に並んでいるグループを全員アトラクションへ誘導する。先頭のグループが待機列から取り出され、アトラクションの空席がグループの人数ぶんだけ減少する。 - 1 に戻る。
ただし、誘導を開始したあとに追加でグループが並ぶことはないとします。以上の条件のもとで、この手順が有限回で終了することが示せます。
高橋君が誘導を開始してから誘導を終了するまで、何回アトラクションをスタートさせるか求めてください。
考察
何の工夫もなく、先頭グループから乗せられるだけアトラクションに乗せていく操作をするだけなので、それを実装すればよいです。実装が少し重い問題ですが、個人的には今回の場合dequeを使うと直感的に書けて楽でした。
コード
from collections import deque
N, K = map(int, input().split())
A = list(map(int, input().split()))
q = deque(A)
cnt = 0
nokori = K
while q:
cnt += 1
while q:
sentou = q.popleft()
if nokori < sentou:
nokori = K
q.appendleft(sentou)
break
else:
nokori -= sentou
print(cnt)
C - Sigma Problem
問題
正整数 $x,y$ に対して $f(x,y)$ を「$(x+y)$ を $10^8$ で割ったあまり」として定義します。
長さ $N$ の正整数列 $A=(A_1,\ldots,A_N)$ が与えられます。次の式の値を求めてください。
$$
\displaystyle \sum_{i=1}^{N-1}\sum_{j=i+1}^N f(A_i,A_j)
$$
考察
足し合わせたとき$10^8$以上となるときが厄介なので、超えるペアと超えないペアとで場合分けすると楽です。そのために、$A$をソートしておいても答えは変わらないので、ソートしておきます。あとは、足し合わせたときに$10^8$以上となる境目をbisect(二分探索)等を使って求め、そのペアの個数分、和から$10^8$引けばよいです。今回の私の実装では2度同じペアを足しわせてしまうので、最後に2で割っています。
コード
import bisect
MOD = 10**8
N = int(input())
A = list(map(int, input().split()))
A.sort()
all = sum(A)
ans = 0
for i, a in enumerate(A):
idx = bisect.bisect_left(A, MOD - a)
k1 = idx
k2 = N - idx
if idx > i:
k1 -= 1
else:
k2 -= 1
tmp = all - a + a * (N - 1) - MOD * k2
ans += tmp
print(ans // 2)
D - Another Sigma Problem
問題
正整数 $x,y$ に対して $f(x,y)$ を以下で定義します。
十進表記の $x,y$ をそれぞれ文字列として解釈しこの順に連結して得られる文字列を $z$ とする。$z$ を十進表記の整数として解釈したときの値を $f(x,y)$ とする。
例えば $f(3,14)=314, f(100,1)=1001$ です。
長さ $N$ の正整数列 $A=(A_1,\ldots,A_N)$ が与えられます。次の式の値を $998244353$ で割ったあまりを求めてください。
$$
\displaystyle \sum_{i=1}^{N-1}\sum_{j=i+1}^N f(A_i,A_j)
$$
考察
例えば$[3,14,15]$の3について注目したとき、314, 315が$f$によって作成されますが、この和は、$300\times2+14+15 = 3 \times 10^2 \times2+14+15$と見ることができます。ここで、14+15は3より右側の数字の総和であり、10の累乗部分は右側の数字の桁数、$\times 2$はその桁数の数字の個数を示しています。よって、この問題を求めるには、左から走査していき、今見ている箇所より右側の数字の総和と、桁数とその個数を管理すればよいです。桁数とその個数の管理にはdefaultdictを用いました。
コード
from collections import defaultdict
MOD = 998244353
N = int(input())
A = list(map(int, input().split()))
nokori = sum(A)
d = defaultdict(int)
for a in A:
d[len(str(a))] += 1
ans = 0
for a in A:
nokori -= a
d[len(str(a))] -= 1
ans += nokori
for key, val in d.items():
ans += a * 10**key * val
ans %= MOD
print(ans)
E - Yet Another Sigma Problem
問題
文字列 $x,y$ に対して $f(x,y)$ を以下で定義します。
- $x,y$ の最長共通接頭辞の長さを $f(x,y)$ とする。
英小文字からなる $N$ 個の文字列 $(S_1,\ldots,S_N)$ が与えられます。次の式の値を求めてください。
$$
\displaystyle \sum_{i=1}^{N-1}\sum_{j=i+1}^N f(S_i,S_j)
$$
考察
先頭から順に見ていって、一致したグループがあれば次の文字について見ていくといった風に探索するのが直感的なので、それを再帰で実装します。このとき、特に今回の実装でいうとdfsの引数s(今のところ先頭からすべて文字が一致している文字列のインデックス群)の部分の計算量が気になるかもしれませんが、これはすべてのノードを合わせても要素数が$\sum_{i=1}^{N} |S_i|$にしかならず、これは$3 \times 10^5$以下なので問題ないです。
コード
from types import GeneratorType
from collections import defaultdict
def bootstrap(f, stack=[]): # yield
def wrappedfunc(*args, **kwargs):
if stack:
return f(*args, **kwargs)
else:
to = f(*args, **kwargs)
while True:
if type(to) is GeneratorType:
stack.append(to)
to = next(to)
else:
stack.pop()
if not stack:
break
to = stack[-1].send(to)
return to
return wrappedfunc
N = int(input())
S = list(map(str, input().split()))
ans = 0
@bootstrap
def dfs(s, idx):
global ans
d = defaultdict(list)
for i in s:
if len(S[i]) > idx:
d[S[i][idx]].append(i)
for val in d.values():
if len(val) >= 2:
ans += len(val) * (len(val) - 1) // 2
yield dfs(val, idx + 1)
yield None
dfs([i for i in range(N)], 0)
print(ans)