mean_abs_error’s blog

ソフトウェアとシステムの狭間

【AtCoder】PythonでABC320を解く

はじめに

ABC320のA,B,C,D,E問題をPython3で解きます。

A問題: Leyland Number

atcoder.jp

問題の概要

A^B + B^Aを計算せよ。

問題の考察

pythonだとべき乗はpowで計算できます。

コード

def main():
    a, b = map(int, input().split())
    ans = pow(a, b) + pow(b, a)
    print(ans)


if __name__ == "__main__":
    main()

B問題: Longest Palindrome

atcoder.jp

問題の概要

文字列 Sについて、 Sの連続する部分文字列のうち、回文であるものの長さの最大値を求めよ。

問題の考察

|S| \le 100とのことなので、全探索できそうです。連続する部分文字列の先頭と長さを全探索するのに \mathcal{O} (|S|^2)、その文字列が回文かを判定するのに \mathcal{O} (|S|)だけ掛かるため、全探索と回文判定で \mathcal{O} (|S|^3)の計算量になります。

コード

def solve(S):
    ans = 1
    for i in range(len(S)):
        for j in range(i + 1, len(S) + 1):
            subS = S[i:j]
            if subS == subS[::-1]:
                ans = max(ans, len(subS))

    return ans


def main():
    S = input()
    ans = solve(S)
    print(ans)


if __name__ == "__main__":
    main()

C問題: Slot Strategy 2 (Easy)

atcoder.jp

問題の概要

N個のリールからなるスロットがあり、 i番目のリールは数字のみからなる長さ Mの文字列 S_iで表される。リールのスロットを止めるスイッチについて押す押さないを毎秒選べ、 t秒目に i番目のリール用のスイッチを押すと S_iの (t \mod M) + 1文字目で止まる。 N=3のとき、3個のリールとも同じ数字になるよう止めるには最小で何秒だけ掛かるか。

問題の考察

10種類( \omega=10)ある0～9の各数字で揃えようとしたときのそれぞれの最小時間を求め、さらにそれらの最小時間の最小を取る方向で考えます。揃えようとする数字を決めたとき、リールを止める順序は N !個だけあり、ある順序で揃えられる場合は N M秒未満で揃えられます。結果として、 \mathcal{O}(\omega N M N!)の計算量で最小時間を求められました。

コード

from itertools import permutations


def solve(M, Ss):
    def measure_time(order):
        i = 0
        for t in range(len(Ss) * M):
            if Ss[order[i]][t % M] == target:
                i += 1
                if i >= len(order):
                    return t

        return -1

    ans = float("inf")
    for target in range(10):
        for order in permutations(range(len(Ss)), len(Ss)):
            t = measure_time(order)
            if t >= 0:
                ans = min(ans, t)

    if ans == float("inf"):
        ans = -1

    return ans


def main():
    M = int(input())
    Ss = [list(map(int, list(input()))) for _ in range(3)]
    ans = solve(M, Ss)
    print(ans)


if __name__ == "__main__":
    main()

D問題: Relative Position

atcoder.jp

問題の概要

座標平面上に N人おり、『人 A_mから見て人 B_mは、 x軸正方向に X_m、 y軸正方向に Y_mだけ離れた位置にいる』という関係性が M個だけ与えられる。人1が原点にいると分かっているとき、それぞれの人がいる座標を求めよ。

問題の考察

位置関係を有向グラフで表現し、人1から幅優先探索(BFS)すればよさそうです。BFSするため、 \mathcal{O} (N + M) の計算量で求められます。

コード

from collections import deque


def solve(N, M, A, B, X, Y):
    G = [[] for _ in range(N)]
    for u, v, x, y in zip(A, B, X, Y):
        G[u].append((v, x, y))
        G[v].append((u, -x, -y))

    P = [None for _ in range(N)]
    P[0] = (0, 0)
    fifo = deque([0])
    used = set([0])
    while fifo:
        u = fifo.popleft()
        for v, x, y in G[u]:
            if v not in used:
                P[v] = (P[u][0] + x, P[u][1] + y)
                fifo.append(v)
                used.add(v)

    return P


def main():
    N, M = map(int, input().split())
    A = []
    B = []
    X = []
    Y = []
    for _ in range(M):
        a, b, x, y = map(int, input().split())
        a -= 1
        b -= 1
        A.append(a)
        B.append(b)
        X.append(x)
        Y.append(y)
    P = solve(N, M, A, B, X, Y)
    for pos in P:
        if pos is None:
            print("undecidable")
        else:
            print(*pos)


if __name__ == "__main__":
    main()

E問題: Somen Nagashi

atcoder.jp

問題の概要

そうめん流しの会に、列の先頭から順に1～Nの番号が付いた N人が来ました。『①時刻 T_mに流れる量 W_mのそうめんを列の先頭にいる人が全て食べて列から外れ、②時刻 T_m + S_mにその人が列の元の位置に戻ってくる』という出来事が M回起こるとき、それぞれの人が食べるそうめんの合計を求めよ。

問題の考察

T_mも S_mも最大で 10^9と大きいため、毎秒シミュレーションするのは計算量的にムリだと分かります。

①そうめんが流れるイベントと②人が戻ってくるイベントはそれぞれ M回しか起こらないため、イベントの発生時刻のみをシミュレーションする方向で考えます。イベントを順に取り出したり、列の先頭にいる人を取り出したりするために、①と②のイベントをヒープで、同様に列に並んだ人もヒープで管理します。イベント用ヒープに挿入するイベントについて、①のイベントを (t, 1, -1, w, s)で、②の人iが戻るイベントを (t, 0, i, -1, -1)で表すことで、適切な順でヒープからイベントを取り出せるようになります。2つのヒープを用いてシミュレーションしつつ、各人の食べたそうめんの量を計算すると、 \mathcal{O} (M \log M + M \log N + N)の解けました。

コード

from heapq import heappop, heappush


def solve(N, M, events):
    ans = [0] * N

    customers = []
    for i in range(N):
        heappush(customers, i)

    while events:
        t, event_id, i, w, s = heappop(events)
        if event_id == 0:
            heappush(customers, i)
        else:
            if customers:
                i = heappop(customers)
                ans[i] += w
                heappush(events, (t + s, 0, i, -1, -1))

    return ans


def main():
    N, M = map(int, input().split())
    events = []
    for _ in range(M):
        t, w, s = map(int, input().split())
        heappush(events, (t, 1, -1, w, s))
    ans = solve(N, M, events)
    print(*ans)


if __name__ == "__main__":
    main()