🍄

Atcoder - ABC340 A - E まで復習解説

2024/02/11に公開

競技プログラミング

AtCoder

tech

ABC340 A - E を復習してまとめてみました。初学者の方の参考になれば幸いです。

A - Arithmetic Progression

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

Python の range 関数がちょうどいいです。 range(A, B + 1, D) とすることで問題と同じ数列を出力することができます。 range は開区間なので range(A, B, D) とすると $B$ が含まれませんので $B + 1$ としましょう。

A, B, D = map(int, input().split())
ans = list(range(A, B + 1, D))
print(*ans)

B - Append

こちらも Python のリストを用いると簡単に実装できます。まさに Python のリストの末尾に $x$ を追加することを append といいますね。

末尾から $k$ 番目の値を求めるのは Python の場合は $A[-x]$ とすれば表示できます。

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

C - Divide and Divide

「 $N$ を $2$ で割って切り捨てた値と切り上げた値にするという操作を繰り返しすべての値を $1$ にするのに必要な金額はいくらですか？」という問題です。切り捨てるのと切り上げるのでほぼ $N / 2$ の値が 2 つ新たにできますし、 $N$ が 2 で割り切れる場合は同じ値がで新たにできます。ほぼ同じ作業を無駄に何度も繰り返すと TLE になってしまいます。

そこで、メモ化再帰をしましょう。

$f(x) := 値 x から生成される値をすべて 1 にするのに必要な金額の総和$

と、定義します。 Python は @cache デコレータをつけると自動で $f(2)$ , $f(10)$ などの値を実行した時に保存してくれるのでその後 $f(2)$ , $f(10)$ を実行したときは $O(1)$ で値を取得できます。求める答えは $f(N)$ です。

$2 \lt N \lt 10^{17}$ と $N$ の値が非常に大きいですが $N$ の値はどんどん $2$ で割られていき、 $log2(10^{17}) = 56.4$ より $1$ にするのに $57$ 回しかかかりません。 $N$ を $2$ で割って生成した値はさらに $57$ 回より少ないのでメモ化再帰を行っておくと TLE にはなりません。

from functools import cache

@cache
def f(x):
    # 整数 x を 処理するのに払う金額
    if x == 1:
        return 0
    a = (x + 1) // 2
    b = x // 2
    return f(a) + f(b) + x

N = int(input())
ans = f(N)
print(ans)

D - Super Takahashi Bros.

各ステージからの移動経路が２パターンあるのですべての経路を調べていると計算量が $O(2^{N})$ (もっと？) となり TLE となってしまいます。

そこで DP かな？と思ったのですが、厄介なのが $1 \le i \le X_{i} \le N$ ではなく $1 \le X_{i} \le N$ であり、前のステージに戻るルートがあります。
(スーパーマリオブラザーズ２みたいに前に戻ってしまうワープがあるイメージでしょうか)

これでは単純に前から DP で値を更新することができません。例えば、本当は $1$ → $9$ → $3$ → $N$ というルートが最短の時に $1$ → $2$ → $3$ → ... → $N$ と更新して行くと $3$ に関する正しい値が得られなくなってしまいです。

そこで経路探索とみなしてダイクストラ法で解きましょう。ダイクストラ法だと heap queue を用いることで、先に更新すべき値から順に計算を行ってくれます。２つの行動はどちらも

ステージ $i$ からステージ $i + 1$ へ $A_{i}$ 秒かけて移動
ステージ $i$ からステージ $X_{i}$ へ $B_{i}$ 秒かけて移動

という経路とみなすことができます。

こうしてみるとダイクストラ法って自動でトポロジカルソートしている DP のようなものとみなすこともできますよね。奥が深いです。

import heapq

N = int(input())
G = [[] for _ in range(N)]

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


def dijkstra(s):
    # 始点sから各頂点への最短距離
    d = [float("inf")] * N
    used = [False] * N
    d[s] = 0
    used[s] = True
    queue = []
    for e in G[s]:
        heapq.heappush(queue, e)
    while queue:
        min_edge = heapq.heappop(queue)
        # まだ使われてない頂点の中から最小の距離のものを探す
        if used[min_edge[1]]:
            continue
        v = min_edge[1]
        d[v] = min_edge[0]
        used[v] = True
        for e in G[v]:
            if not used[e[1]]:
                heapq.heappush(queue, (e[0] + d[v], e[1]))
    return d

dist = dijkstra(0)
ans = dist[-1]
print(ans)

E - Mancala 2

工夫しつつシミュレーションを実際に行うのが早そうです。環状になっているのが少しややこしいですが、箱 $B_{i}$ のボール $x$ 個を $N$ 個の箱に配ると考えたら $x$ を $N$ で割った商を $y$ 余りを $z$ としたら、 $N$ 個の箱に $y$ 個分配して余りは $1$ つずつ箱 $B_{i + 1}$ , $B_{i + 2}$ , ..., $B_{i + z}$ に配ればよいです。途中で $i + z$ が $N$ を超えたら $B_{0}$ , $B_{1}$ , ... としましょう。

この時、愚直に順番に各箱に配ってしまうと $O(N)$ かかりますが区間で配っているので「区間 $l$ から $r$ の箱に $i$ 個配る」という作業をまとめて高速に行うことができればよさそうです。最初は累積和を取ればいいかな？と思ったのですが累積和だと途中で箱 $B_{i}$ のボールの個数を確認することができません。

そこで Lazy Segment Tree を用います。 Lazy Segment Tree は一点更新や区間更新を高速に行うことができます。 Python の AtCoder 環境は公式の C++ ライブラリを Python に移植していただいている https://github.com/not522/ac-library-python がインストールされていますのでそのまま使えます。

from atcoder.lazysegtree import LazySegTree

def op(a, b):
    return 0

e = 0

def mapping(a, b):
    return a + b

def composition(a, b):
    return a + b

id_ = 0

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

segtree = LazySegTree(op, e, mapping, composition, id_, A)

for b in B:
    balls = segtree.get(b)
    # b から balls 個取り出す
    segtree.set(b, 0)
    # b + 1 から順に1個ずつ配る円形状に並んでいるので N - 1 の次は 0 になるように配る
    div, mod = divmod(balls, N)
    # 全部に div 個配り, b + 1 から b + mod までに1追加する
    if div != 0:
        segtree.apply(0, N, div)
    if b + mod < N:
        segtree.apply(b + 1, b + mod + 1, 1)
    else:
        segtree.apply(b + 1, N, 1)
        segtree.apply(0, ((b + mod) % N) + 1, 1)

print(*[segtree.get(i) for i in range(N)])

なお、私はコンテスト中は op, e, mapping, composition の持ち方がイマイチ (区間和とかでググってもうまく例を見つけることができなかったです) ＆ get method に気づかず値を参照するのに prod(i, i + 1) としていたためか TLE になってしまったので、 ChatGPT に質問しながら C++ で実装し直し AC しました。