mean_abs_error’s blog

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

【AtCoder】PythonでABC317を解く

はじめに

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

A問題: Potions

atcoder.jp

問題の概要

N個の効き目順に並んだ傷薬があり、 i番目の傷薬を使うと体力が P_iだけ増加する。現体力が Hで傷薬を1つ使って体力を X以上するとき、その条件を満たす最も効き目の弱い傷薬の番号を答えよ。

問題の考察

H + P_i \ge Xを満たす最も効き目の弱い傷薬番号 iを答えます。

コード

def solve(H, X, P):
    for i, p in enumerate(P, 1):
        if H + p >= X:
            return i

    return len(P)


def main():
    N, H, X = map(int, input().split())
    P = list(map(int, input().split()))
    ans = solve(H, X, P)
    print(ans)


if __name__ == "__main__":
    main()

B問題: MissingNo.

atcoder.jp

問題の概要

N+1個の連続する整数のうちの一つをなくした。入力として与えられる残りの整数の列 A_1, A_2, \cdots, A_Nから、なくした整数を求めよ。ただし、なくした整数は一意に決まる入力のみが与えられる。

問題の考察

なくした整数は一意に決まる入力のみが与えられるという制約から、両端の整数以外がなくなるのだと分かります。両端の整数は必ず残っているため、集合 [ \min(A), \max(A) ]から A_iを除いていくと、残った1つがなくした整数に対応します。

コード

def solve(N, A):
    minA = min(A)
    maxA = max(A)
    V = set(range(minA, maxA + 1))
    for a in A:
        V.discard(a)

    return V.pop()


def main():
    N = int(input())
    A = list(map(int, input().split()))
    ans = solve(N, A)
    print(ans)


if __name__ == "__main__":
    main()

C問題: Remembering the Days

atcoder.jp

問題の概要

N個の街と M本の道路がある。 i本目の道路は長さ C_iで街 A_iと街 B_iを双方向に結ぶ。好きな街からスタートして同じ街を通らずに移動するとき、通る道路の長さの和の最大値を求めよ。

問題の考察

Nが最大でも10と小さいため、深さ優先探索(DFS)による全探索もしくはbitDPで解けそうです。コンテスト中にはDFSによる方法しか思いつかなかったです。

スタートする街を全て試して同じ街を通らないようDFSすると、計算量は \mathcal{O}(N!)となります。なお、DFSを楽に実装できる再帰関数だとTLEが出てしまったため、再帰関数ではなくオイラーツアーで実装しました。

bitDPだと計算量は \mathcal{O}(M 2^N)となります。計算量的にはこちらの方がよいですが、実装にはやや慣れが必要かもしれません。

DFSで解く場合のコード

def solve(N, M, G):
    def dfs(start):
        lifo = [(~start, 0), (start, 0)]
        costs = [-1] * N
        costs[start] = 0
        used = set()
        while lifo:
            u, c = lifo.pop()
            if u >= 0:
                used.add(u)
                for v, w in G[u]:
                    cc = c + w
                    if v not in used:
                        lifo.append((~v, cc))
                        lifo.append((v, cc))
                        costs[v] = max(costs[v], cc)
            else:
                u = ~u
                used.discard(u)

        return costs

    ans = 0
    for start in range(N):
        costs = dfs(start)
        ans = max(ans, max(costs))

    return ans


def main():
    N, M = map(int, input().split())
    G = [[] for _ in range(N)]
    for _ in range(M):
        u, v, w = map(int, input().split())
        u -= 1
        v -= 1
        G[u].append((v, w))
        G[v].append((u, w))
    ans = solve(N, M, G)
    print(ans)


if __name__ == "__main__":
    main()

bitDPで解く場合のコード

from itertools import chain


def solve(N, M, G):
    dp = [[-float("inf")] * N for _ in range(1 << N)]
    for u in range(N):
        dp[1 << u][u] = 0

    for V in range(1 << N):
        for u in range(N):
            for v, w in G[u]:
                if (V >> v) & 1 == 0:
                    VV = V | (1 << v)
                    dp[VV][v] = max(dp[VV][v], dp[V][u] + w)

    ans = max(chain.from_iterable(dp))
    if ans == float("inf"):
        ans = -1

    return ans


def main():
    N, M = map(int, input().split())
    G = [[] for _ in range(N)]
    for _ in range(M):
        u, v, w = map(int, input().split())
        u -= 1
        v -= 1
        G[u].append((v, w))
        G[v].append((u, w))
    ans = solve(N, M, G)
    print(ans)


if __name__ == "__main__":
    main()

D問題: President

atcoder.jp

問題の概要

N個の選挙区において、 i番目の選挙区には X_i人の高橋派と Y_i人の青木派がおり、多数派がその区の Z_i議席を全て獲得する。選挙区全体として過半数を獲得した方が選挙に勝利するとき、高橋君が勝利するためには青木派から高橋派に最低何人を鞍替えさせる必要があるか。

問題の考察

ナップサック問題と同じ構造の問題であるため、動的計画法(DP)で解けそうです。 i番目の選挙区までをみたときに m議席だけ獲得するのに鞍替えが必要な最少人数を \mathrm{dp} [i] [m]とすれば、状態遷移を表せそうです。

i区で鞍替えを考慮する必要がない X_i > Y_iの場合、 \mathrm{dp} [i+1] [m + Z_i ] = \min( \mathrm{dp} [i+1] [m + Z_i ], \mathrm{dp} [i] [m])と配る緩和式を書けます。

それ意外の場合、 i区で議席を獲得するには S_i = \frac{X_i+Y_i + 1}{2} - X_iだけ鞍替えさせる必要があるため \mathrm{dp} [i+1] [m + Z_i ] = \min( \mathrm{dp} [i+1] [m + Z_i ], \mathrm{dp} [i] [m] + S_i)と、議席を獲得しないなら \mathrm{dp} [i+1] [m ] = \min( \mathrm{dp} [i+1] [m], \mathrm{dp} [i] [m])と、配る緩和式を書けます。

状態として iと mを持つため、 \mathcal{O}(N \sum_i Z_i)の計算量でこの問題は解けました。

コード

考察では状態として iを持たせていますが、なくても解けるため省略して解いています。

def solve(N, X, Y, Z):
    dp = [float("inf")] * 200001
    dp[0] = 0
    for x, y, z in zip(X, Y, Z):
        for m in range(100000, -1, -1):
            if x > y:
                dp[m + z] = min(dp[m + z], dp[m])
            else:
                sw = (x + y) // 2 + 1 - x
                dp[m + z] = min(dp[m + z], dp[m] + sw)

    majority = sum(Z) // 2 + 1
    ans = min(dp[majority:])

    return ans


def main():
    N = int(input())
    X = []
    Y = []
    Z = []
    for _ in range(N):
        x, y, z = map(int, input().split())
        X.append(x)
        Y.append(y)
        Z.append(z)
    ans = solve(N, X, Y, Z)
    print(ans)


if __name__ == "__main__":
    main()