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

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

A - Arithmetic Progression

問題

初項が $A$ 、末項が $B$ 、公差が $D$ であるような等差数列を出力してください。

考察

range(x,y,z) と書くと、$x$ 以上 $y$ 未満 の整数を $z$ とばしで見ることができます。
具体的にはこんな感じ。

"""range(x, y, z) の使用例 """
for i in range(3, 23, 4):
    print(i)
# 出力結果
# 3
# 7
# 11
# 15
# 19

上の例では $3$ 以上 $23$ 未満 の整数を $4$ とばしで出力しました。

これを利用すると正解コードがつくれます。
この問題では $A$ 以上 $B$ 以下 の整数を $D$ とばしで出力することに注意してください。

コード

A, B, D = map(int, input().split())

for i in range(A, B + 1, D):
    print(i, end=' ')

別解

リスト内包表記 + リスト出力

この問題ではわざわざリスト表記を持ち出さなくても $3$ 行で正解できましたが、リスト内包表記を知っていると簡潔に書けることがそこそこあります。

また、リストの先頭に * をつけて出力すると、各要素を半角スペース区切りで出力できます。

A, B, D = map(int, input().split())

ans_lst = [i for i in range(A, B + 1, D)]
print(*ans_lst)

print() のキーワード引数endについて

普段は `print()` をつかうと、出力した後に強制で改行されています。
これは、Pythonがデフォルトで出力の末尾に改行コード `\n` を入れているからです。
出力の末尾は、キーワード引数endを指定することで自由にカスタマイズできます。

出力結果は、ぜひ皆さんの手元のPythonで実行して確認してみてください。

"""print()にキーワード引数endを指定する例"""
# どちらも末尾に改行コード\n が入っていて、出力した後に改行される。
print(3)
print(7, end='\n')

print("======================")

# 改行しないとこうなる。
print("At", end='')
print("Cod", end='')
print("er", end='')

print("======================")

# 改行せずに半角スペース区切りの出力をする。
print("We", end=' ')
print("use", end=' ')
print("Python", end=' ')
print("everyday")

B - Append

問題

最初リスト $A$ は空です。
クエリ $Q$ 個を順に処理してください。クエリは以下の $2$ 種類です。

• 1 x : リスト $A$ の末尾に $x$ を挿入する
• 2 k : リスト $A$ の末尾から $k$ 番目の数字を出力する

考察

A.append(x) と書くと、リスト $A$ の末尾に $x$ を挿入できます。
また、リスト $A$ の末尾から $k$ 番目の値は、 A[-k] で取得できます。

コード

Q = int(input())
A = []
for _ in range(Q):
    t, x = map(int, input().split())
    if t == 1:
        A.append(x)
    else:
        print(A[-x])

C - Divide and Divide

問題

黒板に数字 $N$ が $1$ つだけ書かれています。
黒板に書かれている $2$ 以上の数字がなくなるまで、以下の操作を続けます。

• 黒板に書かれている $2$ 以上の数字を $1$ つ選ぶ
• 選んだ数字を $x$ とする。 $x$ を $1$ つ黒板から消して、代わりに $\left \lfloor \dfrac{x}{2} \right\rfloor, \left\lceil \dfrac{x}{2} \right\rceil$ の $2$ つの数字を黒板に書き足す
• この操作をすると $x$ 円のコストがかかる

操作ができなくなるまで操作をし続けたとき、かかった金額の総和はいくらですか？

考察1: 実験してみる

黒板に書かれる数の種類は、実はそんなに多くありません。
これは実験してみると気づけます。
例えばこんなコードを書いて実験してみるといいでしょう。

from sortedcontainers import SortedSet

N = 103  # 適当な整数
ss = SortedSet()
ss.add(N)

lst = []
while ss:
    v = ss.pop()
    lst.append(v)
    for nv in [v // 2, (v + 1) // 2]:
        if nv >= 2: ss.add(nv)

print(len(lst))  # 12

とりあえずPCが固まらないように小さな数で実験してみて、どんどん大きな数にしていくといいです。実際に実験してみると、$N$ に $17$ 桁の整数を入れても大したことない計算量なのが分かるはずです。

考察2: 問題を言い換える

問題では黒板に書かれた数字を $1$ つ消去して $2$ つの数字を新たに書き込んでいました。
これをそのままシミュレーションすると、サンプル3 で出力例が $5655884811924144128$ になっていることから、TLEしそうなのが分かります。

黒板に書かれる数字の種類が多くないことは分かっているので、ここに注目します。

問題での操作を次のように変更します。

• 黒板に書かれた整数の中で最大のものを $x$ とし、その個数を $cnt$ とする
• 黒板に $\left \lfloor \dfrac{x}{2} \right\rfloor, \left\lceil \dfrac{x}{2} \right\rceil$ の $2$ 種類の数をそれぞれ $cnt$ 個ずつ黒板に書く
• この操作でかかるコストは、 $x \times cnt$ 円である

これなら、黒板に書かれる数字の種類数が少ないことから、TLEしなくて済みそうです。

いろいろな実装方法が考えられますが、下ではSortedDictを使用したコードを書いています。

コード

from sortedcontainers import SortedDict

N = int(input())

dic = SortedDict()
dic[N] = 1

ans = 0
while dic:
    num, cnt = dic.popitem()
    ans += num * cnt
    for v in [num // 2, (num + 1) // 2]:
        if v == 1:
            continue
        if v not in dic:
            dic[v] = 0
        dic[v] += cnt

print(ans)

D - Super Takahashi Bros.

問題

$i=1, 2, \cdots , N$ の番号がついた $N$ このステージがあります。最初遊べるのはステージ $1$ だけです。
$i=1, 2, \cdots ,N-1$ についてステージ $i$ が遊べるとき、次の　$2$ 種類のいずれかを行えます。

• $A_i$ 秒かけてステージ $i+1$ を遊べるようにする
• $B_i$ 秒かけてステージ $X_i$ を遊べるようにする

ステージ $N$ が遊べるようになるまでの最短時間を求めてください。

考察

$$dp[s]: \text{ステージsで遊べるようにするための最短時間}$$
とします。

最初はステージ $1$ しか遊べるステージがないですが、ターン数が増えるにつれて遊べるステージの候補も増えてきます。
次のようにステージを選んでいくことにします。

• まだ選ばれていないステージの中で、 $dp[s]$ の値が最も小さいステージを選ぶ。このステージを $i$ とする
• $dp[i+1], dp[X_i]$ の $2$ つの値を必要があれば更新する

この実装は下のコードを参考にしてください。

また、これは「ダイクストラ法」とよばれる典型アルゴリズムです。この機会におぼえておきましょう。

コード

from heapq import heappush, heappop

INF = 1 << 63

N = int(input())
G = []
for _ in range(N - 1):
    a, b, x = map(int, input().split())
    G.append((a, b, x - 1))

# dp[s]:ステージsで遊べるようにするための最短時間
dp = [INF] * N
dp[0] = 0
que = []
heappush(que, (dp[0], 0))
while que:
    time, stage = heappop(que)
    if dp[stage] != time:
        continue
    if stage + 1 >= N:
        continue
    a, b, x = G[stage]
    if dp[stage + 1] > (nex_time := dp[stage] + a):
        dp[stage + 1] = nex_time
        heappush(que, (dp[stage + 1], stage + 1))
    if dp[x] > (nex_time := dp[stage] + b):
        dp[x] = nex_time
        heappush(que, (dp[x], x))

print(dp[N - 1])

E - Mancala 2

問題

(少し複雑なので、問題概要は省略します)

考察

実際にいくつかシミュレーションしてみると、どの箱を見ても配られるボールの数はだいたい同じで、差は高々 $1$ つだけです。
ある区間を指定して、その範囲全体の数字をいくつかプラスできると嬉しいです。
これは、遅延セグメントツリーをつかうと実現できます。

遅延セグメントツリーのつかい方は、下の記事を読んでください。

• Pythonで遅延セグメントツリーの問題を解けるようにする！
  • 遅延セグメントツリーをまったく知らない人向け。「遅延セグメントツリーが想定解法の簡単な問題が解けるようになる」がこの記事の目標です
• ACL Practice Contestをac-library-pythonで解いてみたよ。(G～L問題)
  • ac-library-pythonにある遅延セグメントツリーのメソッドをすべて列挙し、そのつかい方を書いています

コード

from atcoder.lazysegtree import LazySegTree

e = 0


# 区間演算
def op(ele1, ele2):
    return ele1 + ele2


def mapping(func, ele):
    return func + ele


def composition(func_upper, func_lower):
    return func_upper + func_lower


# mapping(func, ele) = ele となるようなfunc
id_ = 0

N, M = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
lazy_segtree = LazySegTree(op, e, mapping, composition, id_, A)
for b in B:
    ball_cnt = lazy_segtree.get(b)
    lazy_segtree.set(b, 0)
    ave = ball_cnt // N  # 全員に配るぶん
    if ave:
        lazy_segtree.apply(0, N, ave)
    zan = ball_cnt - ave * N
    if b + zan < N:
        lazy_segtree.apply(b + 1, b + 1 + zan, 1)
    else:
        if b + 1 < N:
            lazy_segtree.apply(b + 1, N, 1)
        zan -= N - b - 1
        lazy_segtree.apply(0, zan, 1)

ans_lst = [lazy_segtree.get(i) for i in range(N)]
print(*ans_lst)