ABC353 (A~E) をPythonで解く

A問題：Buildings

一番左のビルより高いビルが存在するかを順番に確認し、見つけたら終了。インデックスが1始まりなことに注意する。

コード

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

leftmost = H[0]
ans = -1

for i in range(1, N):
    if H[i] > leftmost:
        ans = i+1
        break

print(ans)

B問題：AtCoder Amusement Park

列に並んでいる人を順番に確認する。ここで、アトラクションの残りの席数をavailable_seatsとして考える。

もし先頭のグループの人数がavailable_seatsよりも多かった場合、そのグループは案内できないのでアトラクションを出発させ、空席を$K$に戻す。そしてそのグループをアトラクションに案内し、空席を人数分減らす。

もし先頭のグループの人数がavailable_seatsよりも少なかった場合、そのグループを案内して空席を人数分減らす。

どちらの場合にせよ空席を人数分減らすので、そこはまとめて処理する。

全員分確認し終わった後、もし最後のアトラクションがまだ出発していなければ、出発させる。これは空席が$K$かどうかで判別できる。空席が$K$でないということはもう乗車していて出発を待っている人がいるということなので。

コード

N, K = map(int, input().split())
from collections import deque
A = deque(list(map(int, input().split())))

ans = 0
available_seats = K

for people in A:
    if people > available_seats:
        ans += 1
        available_seats = K
    available_seats -= people
    
if available_seats != K:
    ans += 1

print(ans)

C問題：Sigma Problem

$\Sigma$が出てきてややこしい。$\Sigma$の読み解き方はこの記事などに詳しいが、ざっくりいうと総和を取るための記号。

$$\sum_{i=1}^{N-1} \sum_{j=i+1}^N f(x_i, x_j)$$

これは、

1. $i$を1から$N-1$まで全部確認する。
2. それぞれの$i$において、$j$を$i+1$から$N$まで取る。
3. $f(x_i, x_j)$の値を取る。
4. 全部足し合わせる。

という意味。

とりあえず愚直解

コードで書くと、

for i in range(0, N-1):
    for j in range(i+1, N):
        ans += (A[i] + A[j]) % 10**8

で愚直解が生成できる。

つまり数列内の要素でペアを全部作ってその和を考える。

但し、この足し算では計算量が$O(N^2)$となって間に合わないので、より計算量の少ない解法を考える必要がある。

解法

数列が与えられるが各要素の順番はそれほど関係がない問題の鉄則として、「とりあえずソートしてみて何かできないか考えてみる」というものがある。ソートされた数列を考え、順番にペアを作って和を計算していくことを考えると、和が単調増加することがわかる。

ということは、いったんペアの和が$10^8$を超えると、そのあとのペアの和は絶対に$10^8$を超える。ただし$A_i<10^8$なので、$A_i+A_j$が$10^8\times2$を超えることはない。

従って、毎回丁寧にペアの和から$10^8$を引いてあげなくとも、和が$10^8＄$を超えるペアの数を数えてその分だけ$10^8$をまとめて引いてあげればよい。

どこから和が$10^8$を超えるかは二分探索を使えば$O(logN)$で計算できる。それを$N$個の要素に対して行うので計算量は$O(NlogN)$でまとめられる。

ペアの和の合計を考える。全ての要素が自分自身以外の要素とペアを作るので、1要素につき$N-1$のペアを作成する。よって、

seq_sum = sum(A) * (N-1)

で計算ができる。

コード

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

from bisect import bisect_left

A.sort()

# 全ペアの和を求めるために、各要素が全ペアに寄与する総和を計算する
total_sum = sum(A) * (N - 1)

# 和が10**8以上のペアのカウント
cnt = 0
for i in range(N):
    # 和が10**8を超える最小のjを探す
    limit = 10**8 - A[i]
    idx = bisect_left(A, limit, i + 1)  # 自分自身を除外するため i + 1 から探す
    cnt += N - idx  # idx以降の要素数が条件を満たす

# 和が10**8を超えるペアの総和がどれだけ余分か計算する
excess = cnt * 10**8

result = total_sum - excess
print(result)

D問題：Another Sigma Problem

考え方

C問題と式は似ているが、操作が変わっている。今回はペアとなった2つの数字を結合して1つの数字と見做し、合計していく問題。

例をもとに考えてみる。

3, 14, 15

結合されて出てくる数字は、それぞれ

• 3 + 14 → 314
• 3 + 15 → 315
• 14 + 15 → 1415

これの合計が答えになる。

3+14を考えると、これは
$3 \times10^{14の桁数=2} + 14$で計算されることがわかる。

3+15も同様に、
$3 \times10^{15の桁数=2} + 15$で計算される。

これ、いい感じに$3\times (10^2+10^2) + 14 + 15$でまとめられそう。
どれくらいの桁数を掛ける必要があるかを逆順の累積和を使って前以て計算しておけば、計算量が減らせる。

具体的には

3, 14, 15

から

210, 200, 100

が作れる。
100は$10^{15の桁数}$
200は$10^{14の桁数}+10^{15の桁数}$
210は$10^{3の桁数}+10^{14の桁数}+10^{15の桁数}$
となっている。

これを使うと、3, 14, 15の組は以下のように計算できる。

• 3を含むペアの3の部分の計算
  • 3 + 何かになるペア → $3\times200=600$
  • 何か + 3になるペア → 3の前に要素が存在しないので0
• 14を含むペアの14の部分の計算
  • 14 + 何かになるペア → $14\times100=1400$
  • 何か + 14になるペア → 14の前に要素が1つあるので$14\times1=14$
• 15を含むペアの15の部分の計算
  • 15 + 何かになるペア → 15の後に要素がないので0
  • 何か + 15になるペア → 15の前に要素が2つあるので$15\times2=30$

全て合計して2200.

あとは 998244353 で割った余りを返すところだけ忘れなければ大丈夫。

解法

桁数の計算をする。

digits = [0] * N  

# 最後の要素の10の桁数べきを設定
digits[-1] = 10 ** len(str(A[-1]))

# 逆順で各要素に対して累積和で処理
for i in range(N-2, -1, -1):
    digits[i] = 10 ** len(str(A[i])) + digits[i+1]

要素を順番に見ていく。

ans = 0 

for j in range(N):
    ans += A[j] * j % MOD # 何か+要素の部分
    # 最後の要素は桁数を考えなくてよい
    if j != N-1:
        ans += A[j] * digits[j+1] % MOD # 要素＋何かの部分
    
    ans %= MOD

print(ans)

コード

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

digits = [0] * N 
digits[-1] = 10 ** len(str(A[-1]))

for i in range(N-2, -1, -1):
    digits[i] = 10 ** len(str(A[i])) + digits[i+1]

ans = 0 

for j in range(N):
    ans += A[j] * j % MOD
    if j != N-1:
        ans += A[j] * digits[j+1] % MOD
    
    ans %= MOD

print(ans)

E問題：Yet Another Sigma Problem

この問題は解説を読んでコードを追ってざっくりとわかっただけなのでまだ解説は描けない。参考になりそうなリンクとコメント付きのコードだけ貼っておく。

公式の解説

Wikiの記事

トライ木がなんですごいのか

参考にした実装例

トライ木自体が初耳の人は、問題の入力例を下記のコードに突っ込んでみてtrie_treeの挙動を追えばざっくりと理解できると思う。簡潔に言うと、辞書のそれぞれのキーの値がそのキーを接頭辞に持つ単語の数とその次に続く文字が辞書として格納されたもののリストになっている。

例えばabを接頭辞として持つ単語の数は、aのキーに内包された辞書の中のbのキーの値を見ればわかる、といった具合。

N = int(input())
S = list(map(str, input().split()))

trie_tree = {}
ans = 0

for word in S:
    # 順次アップデートしているので
    current = trie_tree 

    for letter in word:
        # もうキーとして存在していれば
        if letter in current: 
            # 値をansに追加
            ans += current[letter][0]
            current[letter][0] += 1
        # キーとして存在していなければ
        else:
            # キーを追加する
            current[letter] = [1, {}]

        # 子ノードに移動する
        current = current[letter][1]

print(ans)