2023-11-13

【AtCoder】PythonでABC328を解く

はじめに

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

はじめに
A問題
B問題
C問題
D問題
E問題

A問題

atcoder.jp

問題の概要

配点が $X$ 以下である問題すべての配点の合計を出力せよ。

問題の考察

リスト内包表記とsum関数を使うと楽に実装できます。

コード

def main():
    N, X = map(int, input().split())
    S = list(map(int, input().split()))
    ans = sum(s for s in S if s <= X)
    print(ans)


if __name__ == "__main__":
    main()

B問題

atcoder.jp

問題の概要

$i$ 月 $j$ 日 ( $1 \le i \le N, 1 \le j \le D_i$ )がゾロ目になる日付の日数を数えよ。

問題の考察

日付を全探索し、各日付を文字列に直して集合のサイズが1か調べることでゾロ目を判定します。

コード

def solve(N, D):
    repdigit_dates = []
    for i, d in enumerate(D, 1):
        for j in range(1, d + 1):
            date = f"{i}{j}"
            if len(set(date)) == 1:
                repdigit_dates.append(date)

    return repdigit_dates


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


if __name__ == "__main__":
    main()

C問題

atcoder.jp

問題の概要

文字列 $S$ が与えられる。 $S$ の $l_q$ 文字目から $r_q$ 文字目までからなる部分文字列において、同じ文字が2つ隣り合う箇所はいくつあるか、という $Q$ 個の質問に答えよ。

問題の考察

2つ隣り合う箇所の累積和を事前に求めておくことで、各質問に $\mathcal O (1)$ で答えます。

コード

from itertools import pairwise


def solve(N, Q, S, LR):
    X = [0, 0]
    for s1, s2 in pairwise(S):
        X.append(X[-1] + int(s1 == s2))

    ans = []
    for left, right in LR:
        z = X[right] - X[left]
        ans.append(z)

    return ans


def main():
    N, Q = map(int, input().split())
    S = input()
    LR = [list(map(int, input().split())) for _ in range(Q)]
    ans = solve(N, Q, S, LR)
    print(*ans)


if __name__ == "__main__":
    main()

D問題

atcoder.jp

問題の概要

A,B,Cの3種類の文字からなる文字列 $S$ が与えられる。連続な部分文字列としてABCを含む限り $S$ から削除するというのを繰り返すとき、最終的な $S$ を出力せよ。

問題の考察

$S$ の先頭の文字から順にスタックに積んでいき、スタックの後ろ3文字がABCだった場合はスタックから3文字分だけpopします。Pythonの場合、スライスを使って後ろ3文字を調べることで、現在のスタックのサイズが3未満だったとしてエラーなく判定できます。

コード

def solve(S):
    lifo = []
    for s in S:
        lifo.append(s)
        S3 = "".join(lifo[-3:])
        if S3 == "ABC":
            lifo.pop()
            lifo.pop()
            lifo.pop()

    return "".join(lifo)


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


if __name__ == "__main__":
    main()

E問題

atcoder.jp

問題の概要

$N$ 頂点 $M$ 辺の重み付き単純連結無向グラフと正整数 $K$ が与えられる。このグラフの全域木に対して、含まれる辺の重みの総和を $K$ で割った余りをその全域木のコストとするとき、このグラフの全域木のコストの最小値を求めよ。

問題の考察

制約より最大で頂点数が8で辺数が28であることから、このグラフから作れる全域木の数は高々 ${}_{28}C_{7}$ と分かるため、全探索しても間に合いそうです。辺の集合から $N - 1$ 個の辺を取り出す組合せを全探索します。各組合せでは、UnionFindを使ってサイズが $N$ であるかを調べることで、全域木になっているかを判定します。

コード

from itertools import combinations

from atcoder.dsu import DSU


def solve(N, M, K, edges):
    ans = float("inf")
    for sub_edges in combinations(edges, N - 1):
        uf = DSU(N)
        total = 0
        for u, v, w in sub_edges:
            uf.merge(u, v)
            total += w
            total %= K
        if uf.size(0) == N:
            ans = min(ans, total)

    return ans


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


if __name__ == "__main__":
    main()

2023-10-29

【AtCoder】PythonでABC326を解く

はじめに

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

はじめに
A問題: 2UP3DOWN
B問題: 326-like Numbers
C問題: Peak
E問題: Revenge of "The Salary of AtCoder Inc."

A問題: 2UP3DOWN

atcoder.jp

問題の概要

ビルの $X$ 階から $Y$ 階へ移動するとき、2階分までの上り、または、3階分までの下りであれば階段を使い、そうでないときはエレベータを使う。階段で移動するか判定せよ。

問題の考察

$Y - X$ が範囲内に収まっているかで判定します。

コード

def main():
    x, y = map(int, input().split())
    ans = "Yes" if -3 <= (y - x) <= 2 else "No"
    print(ans)


if __name__ == "__main__":
    main()

B問題: 326-like Numbers

atcoder.jp

問題の概要

3桁の整数について、百の位のかずと十の位の数の積が一の位の数に等しいものを326-like numberと呼ぶ。整数 $N$ 以上の最小の326-like numberを求めよ。

問題の考察

制約を見ると分かるように326-like numberの最大値は919ですので、 $[N, 919$ ]の範囲で順に探索します。各位の数を取り出すには整数を文字列に変換すると楽です。

コード

def solve(N):
    for i in range(N, 920):
        a, b, c = map(int, list(str(i)))
        if a * b == c:
            return i

    return -1


def main():
    N = int(input())
    ans = solve(N)
    print(ans)


if __name__ == "__main__":
    main()

C問題: Peak

atcoder.jp

問題の概要

数直線上の座標 $A_n$ にプレゼントが置かれている。ある座標 $x$ を選ぶと半開区間 $[x, x + M)$ にあるプレゼントを全て獲得できる。適切に $x$ を選ぶと最大でいくつのプレゼントを獲得できるか。

問題の考察

配列 $A$ をソートした上で、 $x$ を各 $A_n$ で全探索しつつ高速に獲得できるプレゼント数をカウントします。二分探索で $x+M$ 未満の数と $x$ 未満の数をそれぞれ求めて差を計算すると、半開区間にあるプレゼント数を $\mathcal O (\log N)$ でカウントできます。前処理としてのソート、および、全探索と二分探索により、 $\mathcal O (N \log N)$ の計算量でプレゼント数の最大値を計算できました。

コード

from bisect import bisect_left


def solve(N, M, A):
    ans = 0
    for a in A:
        x = bisect_left(A, a + M) - bisect_left(A, a)
        ans = max(ans, x)

    return ans


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


if __name__ == "__main__":
    main()

E問題: Revenge of "The Salary of AtCoder Inc."

atcoder.jp

問題の概要

そのままですので、略します。

問題の考察

$f(x)$ をダイスで $x$ が出ているときの給料の期待値とします。ダイスで $x$ が出ているときにさらに手順を踏むと、次のダイス $z$ が $1$ から $x$ までのときはそれぞれ $\frac{1}{N}$ の確率で $A_x$ が支給され、 $x+1$ から $N$ までのときはそれぞれ $\frac{1}{N}$ の確率で $A_x + f(z)$ が支給されます。
$f(x) = \frac{1}{N} \sum_{z=1}^x A_x + \frac{1}{N} \sum_{z=x+1}^N [ A_x + f(z) ] = A_x + \frac{1}{N} \sum_{z=x+1}^N f(z)$

上記の式について、 $x=N$ の状態から $x=0$ の状態まで計算していけば、 $x=0$ のときの給料の期待値が計算できます。なお、高速化のために式の最後の総和では累積和を用います。これによって、計算量 $\mathcal O (N)$ で期待値が求まります。なお、答えを $\bmod$ で計算する必要があるため、 $N$ での割り算は $\mathrm{pow} (N, -1, 998244353)$ の掛け算で実現します。

コード

def solve(N, A, MOD=998244353):
    A = [0] + A
    dp = [0] * (N + 1)
    ep = [0] * (N + 2)
    modinv_N = pow(N, -1, MOD)
    for x in range(N, -1, -1):
        dp[x] = A[x]
        dp[x] += ep[x + 1] * modinv_N % MOD
        dp[x] %= MOD
        ep[x] = dp[x] + ep[x + 1]
        ep[x] %= MOD

    ans = dp[0]

    return ans


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


if __name__ == "__main__":
    main()

2023-10-22

【AtCoder】PythonでABC325を解く

はじめに

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

はじめに
A問題: Takahashi san
B問題: World Meeting
C問題: Sensors
D問題: Printing Machine

A問題: Takahashi san

atcoder.jp

問題の概要

与えられた名字と名前に対して、名字の方に敬称(san)を付けよ。

問題の考察

Pythonで文字列のフォーマットをいじるのに便利なf-stringを使って敬称を付けます。

コード

def main():
    S, T = input().split()
    ans = f"{S} san"
    print(ans)


if __name__ == "__main__":
    main()

B問題: World Meeting

atcoder.jp

問題の概要

ある国の拠点 $n$ の時刻は世界標準時で0時のときに $X_n$ 時で、その拠点には $W_n$ 人の社員が所属している。各拠点の社員は、現地時間で09:00-18:00の時間帯のときのみ会議に参加できる。適切な開始時刻を選んで1時間の会議を設定するとき、会議の参加人数は最大で何人になるか。

問題の考察

世界標準時について0時から23時までを全探索します。ある世界標準時における参加人数は、その世界標準時を拠点 $n$ の時間に直したときに9以上18未満を満たす $W_n$ の総和となります。

コード

def solve(N, W, X):
    ans = 0
    for t in range(24):
        total = sum(w for w, x in zip(W, X) if 9 <= (t + x) % 24 < 18)
        ans = max(ans, total)

    return ans


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


if __name__ == "__main__":
    main()

C問題: Sensors

atcoder.jp

問題の概要

$H \times W$ のマス目にセンサが配置されている。上下左右斜めで隣接しているセンサ同士は、連動して動作する一纏まりのセンサ群として見なせる。センサ群は何個あるか。

問題の考察

幅優先探索やUnionFindで解くことができますが、ここではUnionFindで解きます。上下左右斜めで隣接しているセンサをUnionFindで同じ集合として結合し、最後に各集合の親がセンサか否かでセンサ群を数えます。

UnionFindのソースコードは、こちらから引用しました(https://note.nkmk.me/python-union-find/)。いつもお世話になっております。

コード

from collections import defaultdict


class UnionFind:
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def members(self, x):
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

    def all_group_members(self):
        group_members = defaultdict(list)
        for member in range(self.n):
            group_members[self.find(member)].append(member)
        return group_members

    def __str__(self):
        return "\n".join(f"{r}: {m}" for r, m in self.all_group_members().items())


def solve(H, W, S):
    uf = UnionFind((H + 2) * (W + 2))
    for h in range(1, H + 1):
        for w in range(1, W + 1):
            if S[h][w] == "#":
                u = h * (W + 2) + w
                for dh in range(-1, 2):
                    for dw in range(-1, 2):
                        hh = h + dh
                        ww = w + dw
                        if S[hh][ww] == "#":
                            v = hh * (W + 2) + ww
                            uf.union(u, v)

    ans = 0
    for r in uf.roots():
        h = r // (W + 2)
        w = r % (W + 2)
        ans += int(S[h][w] == "#")

    return ans


def main():
    H, W = map(int, input().split())
    S = []
    S.append("." * (W + 2))
    for _ in range(H):
        line = "." + input() + "."
        S.append(line)
    S.append("." * (W + 2))
    ans = solve(H, W, S)
    print(ans)


if __name__ == "__main__":
    main()

D問題: Printing Machine

atcoder.jp

問題の概要

商品 $n$ は、今から $T_n$ 秒後に印字機に入り、その $D_n$ 秒後に印字機から出る。印字機は、中にある商品に対して一瞬で印字することができるが、印字の度に1秒間のチャージ時間を必要とする。印字する対象の商品と印字のタイミングを適切に選んだとき、最大で何個の商品に印字できるか。

問題の考察

印字機に入っている商品の中で最も出が早いものから印字する貪欲法でよさそうですが、 $T_n, D_n$ が大きいため、印字機への入りと出を毎秒シミュレーションする方法だと計算量的に間に合いません。

時間でのシミュレーションではなく、商品でのシミュレーションによる貪欲法を考えます。まず、印字機に入る時刻で管理された商品用のヒープ(in_heap)と、印字機を出る時刻で管理された商品用のヒープ(out_heap)を持ちます。①in_heap先頭の商品と同じ入り時刻 $\tau$ の商品全てをout_heapに入れ、現時刻を $\tau$ に更新します。②out_heapの先頭の商品の出時刻が $\tau$ 以上であるときは印字してチャージし( $\tau := \tau + 1$ )、そうでないときは印字しない(できない)、というのをin_heap先頭の商品の入り時刻が $\tau$ より大きいうちは繰り返します。①と②をin_heapが空になるまで繰り返しせば、 $\mathcal O (N \log N)$ の計算量で解けました。

コードでは、入りと出の時刻が十分大きい商品を番兵としてin_heapに入れており、これによってin_heapが空になることはないために条件分岐がシンプルになります。

コード

from heapq import heappop, heappush


def solve(N, T, D):
    in_heap = []
    for t, d in zip(T, D):
        heappush(in_heap, (t, d))
    heappush(in_heap, (10000000000000000000, 10000000000000000000))
    out_heap = []

    ans = 0
    while len(in_heap) > 1:
        tau = in_heap[0][0]
        while in_heap[0][0] <= tau:
            t, d = heappop(in_heap)
            heappush(out_heap, t + d)
        while out_heap and tau < in_heap[0][0]:
            if tau <= out_heap[0]:
                heappop(out_heap)
                ans += 1
                tau += 1
            else:
                heappop(out_heap)

    return ans


def main():
    N = int(input())
    T = []
    D = []
    for _ in range(N):
        t, d = map(int, input().split())
        T.append(t)
        D.append(d)
    ans = solve(N, T, D)
    print(ans)


if __name__ == "__main__":
    main()

2023-10-15

【AtCoder】PythonでABC324を解く

はじめに

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

はじめに
A問題: Same
B問題: 3-smooth Numbers
C問題: Error Correction
D問題: Square Permutation

A問題: Same

atcoder.jp

問題の概要

与えられる整数が全て等しいか答えよ。

問題の考察

集合のサイズが1かを調べます。

コード

def main():
    N = int(input())
    A = list(map(int, input().split()))
    ans = "Yes" if len(set(A)) == 1 else "No"
    print(ans)


if __name__ == "__main__":
    main()

B問題: 3-smooth Numbers

atcoder.jp

問題の概要

正の整数 $N$ について、 $N = 2^x 3^y$ を満たす整数 $x, y$ が存在するか答えよ。

問題の考察

$N$ を $2, 3$ で割れるだけ割ったときに $1$ となれば、そのような $x, y$ が存在しています。

コード

def solve(N):
    n = N
    while n % 2 == 0:
        n //= 2
    while n % 3 == 0:
        n //= 3

    return n == 1


def main():
    N = int(input())
    ans = "Yes" if solve(N) else "No"
    print(ans)


if __name__ == "__main__":
    main()

C問題: Error Correction

atcoder.jp

問題の概要

$N$ 個の文字列 $S_1, S_2, \cdots, S_N$ が文字列 $T'$ とほぼ等しいかを判定せよ(ほぼ等しいの意味は問題文を参照)。

問題の考察

文字列 $S_n$ と文字列 $T'$ の長さが等しい場合、1つ差のある場合、それ以外の3つに場合分けて考えます。長さが等しい場合、先頭から順に見て違う文字同士の数が1以下ならほぼ等しいと分かります。1つ差の場合、先頭から順に見て違う文字があったらそれ以降が全て同じ文字ならばほぼ等しいと分かります。それ以外の場合、明らかにほぼ等しくなることはないです。 $N$ 個の文字列の入力と先頭から1文字ずつ見ていく処理があるため、計算量は $\mathcal O (N + |T'|+ \sum_{i=1}^N |S_i|)$ で解けました。

コード

def solve(N, TT, Ss):
    def is_valid(S1, S2):
        if len(S1) > len(S2):
            S1, S2 = S2, S1

        if len(S1) == len(S2):
            return sum([s != tt for s, tt in zip(S1, S2)]) <= 1
        elif len(S1) + 1 == len(S2):
            for i, (s1, s2) in enumerate(zip(S1, S2)):
                if s1 != s2:
                    return S1[i:] == S2[i + 1 :]
            return True

        return False

    ans = []
    for i, S in enumerate(Ss, 1):
        if is_valid(S, TT):
            ans.append(i)

    return ans


def main():
    N, TT = input().split()
    N = int(N)
    Ss = [input() for _ in range(N)]
    ans = solve(N, TT, Ss)
    print(len(ans))
    print(*ans)


if __name__ == "__main__":
    main()

D問題: Square Permutation

atcoder.jp

問題の概要

長さ $N$ の文字列 $S$ を並び替えた結果を10進数の整数として解釈したもののうち、平方数となるようなものはいくつあるか答えよ。

問題の考察

$\max(N) = 13$ と小さくはあるものの、並び替えた整数を全て作るのに $\mathcal O (N !)$ の計算量だけ掛かるため、この方針だと解けないと分かります。

Wikipediaの平方数のページにあった平方数の数列を見ると、平方数の数は明らかに少ないと分かるため、先に平方数を列挙し、その後に文字列 $S$ で各平方数を作れるかを判定するのがよさそうです。平方数を列挙するには、 $0$ から $\sqrt{10^N}$ までの整数を二乗すれば得られます。各平方数を作れるかは、各平方数をゼロパディングした $N$ 桁の10進数と見たときに各桁に現れる整数の頻度が、文字列 $S$ と同じか調べれば判定できます。平方数を列挙するのに $\mathcal O ( \sqrt{10^N} )$ 、各平方数を作れるかの判定に $\mathcal(N)$ だけ掛かるため、全体で $\mathcal O ( N \sqrt{10^N} )$ の計算量で解けました。

コード

from collections import Counter
from itertools import count, takewhile


def solve(N, target_digits):
    target_hist = Counter(target_digits)

    MAX = pow(10, N)
    squares = [i * i for i in takewhile(lambda x: x * x < MAX, count())]

    ans = 0
    for s in squares:
        digits = list(f"{s:0{N}}")
        hist = Counter(digits)
        ans += int(hist == target_hist)

    return ans


def main():
    N = int(input())
    S = input()
    ans = solve(N, list(S))
    print(ans)


if __name__ == "__main__":
    main()

2023-10-08

【AtCoder】PythonでABC323を解く

はじめに

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

はじめに
A問題: Weak Beats
B問題: Round-Robin Tournament
C問題: World Tour Finals
D問題: Merge Slimes

A問題: Weak Beats

atcoder.jp

問題の概要

16bitの数値が与えられる。MSB(最上位ビット)から見た2bit目、4bit目、…、16bit目といった偶数目のビット全てが0かを答えよ。

問題の考察

偶数目のビットを1とした0b0101010101010101とのビット演算子の論理積を取ることで判定します。

コード

def main():
    S = int(input(), 2)
    ans = "No" if S & 0b0101010101010101 else "Yes"
    print(ans)


if __name__ == "__main__":
    main()

B問題: Round-Robin Tournament

atcoder.jp

問題の概要

$1$ から $N$ までの番号が付いたプレイヤー $N$ 人での総当たり戦の結果が与えられる。勝利数が多い方が順位は上になり、勝利数が同じ場合はプレイヤー番号が小さい方が順位は上になる。順位が高い順にプレイヤー番号を答えよ。

問題の考察

勝利数とプレイヤー番号のタプルを作ってソートします。

コード

def solve(N, Ss):
    wins = [0] * N
    for i, S in enumerate(Ss):
        for s in S:
            if s == "o":
                wins[i] += 1

    X = [(-w, i) for i, w in enumerate(wins, 1)]
    X = sorted(X)

    ans = [i for _, i in X]

    return ans


def main():
    N = int(input())
    Ss = [input() for _ in range(N)]
    ans = solve(N, Ss)
    print(*ans)


if __name__ == "__main__":
    main()

C問題: World Tour Finals

atcoder.jp

問題の概要

$N$ 人のプレイヤーそれぞれについて、点数 $A_m$ な $M$ 問の問題をある時刻に解いた結果が与えられる。各プレイヤーについて、その時刻の最高点を上回るのに必要な解くべき問題数の最小値を求めよ。

問題の考察

点数の高い順に問題をソートしておき、まだ解いていない問題から加算していくことで、最高点を超えるのに必要な問題数を計算します。

コード

def solve(N, M, A, Ss):
    scores = [
        i + sum(A[j] for j, s in enumerate(S) if s == "o") for i, S in enumerate(Ss, 1)
    ]
    target = max(scores)

    X = [(a, i) for i, a in enumerate(A)]
    X = sorted(X, reverse=True)

    ans = []
    for i, (S, score) in enumerate(zip(Ss, scores)):
        if score == target:
            ans.append(0)
        else:
            n = 0
            for a, j in X:
                if S[j] == "x":
                    score += a
                    n += 1
                    if score > target:
                        ans.append(n)
                        break

    return ans


def main():
    N, M = map(int, input().split())
    A = list(map(int, input().split()))
    S = [input() for _ in range(N)]
    ans = solve(N, M, A, S)
    print(*ans)


if __name__ == "__main__":
    main()

D問題: Merge Slimes

atcoder.jp

問題の概要

$N$ 種類のサイズのスライムがおり、サイズ $S_n$ のスライムは $C_n$ 匹いる。同じサイズのスライムが2匹いるとき、合成することでその2匹は消滅して代わりに倍のサイズのスライム1匹になる。最小でスライムは何匹になるか。

問題の考察

小さいサイズから順に合成していけばよさそうです。ヒープに初期値として $N$ 種類のサイズを入れておき、小さいサイズから取り出しつつ合成したら倍のサイズをヒープに入れます。

コード

from collections import Counter
from heapq import heappop, heappush


def solve(N, S, C):
    cnts = Counter()
    for s, c in zip(S, C):
        cnts[s] += c

    heap = []
    for s in S:
        heappush(heap, s)

    while heap:
        s = heappop(heap)
        if cnts[s] > 1:
            cnts[2 * s] += cnts[s] // 2
            cnts[s] %= 2
            heappush(heap, 2 * s)

    return sum(cnts.values())


def main():
    N = int(input())
    S, C = list(zip(*[list(map(int, input().split())) for _ in range(N)]))
    ans = solve(N, S, C)
    print(ans)


if __name__ == "__main__":
    main()

2023-10-03

【AtCoder】PythonでABC322を解く

はじめに

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

はじめに
A問題: First ABC 2
B問題: Prefix and Suffix
C問題: Festival
D問題: Polyomino

A問題: First ABC 2

atcoder.jp

問題の概要

文字列 $S$ の中で文字列ABCが部分文字列として初めて現れる位置を答えよ。現れない場合は-1を出力せよ。

問題の考察

pythonの文字列クラスにあるfindメソッドを使うと楽です。

コード

def solve(S, TARGET="ABC"):
    return S.find(TARGET)


def main():
    _ = input()
    S = "@" + input()
    ans = solve(S)
    print(ans)


if __name__ == "__main__":
    main()

B問題: Prefix and Suffix

atcoder.jp

問題の概要

文字列 $S, T$ が与えられる。 $S$ が $T$ の接頭語であるかないか、 $S$ が $T$ の接尾語であるかないかで計4パターン存在する。どのパターンかを答えよ。

問題の考察

接頭語か否かの判定結果(is_head)と接尾語か否かの判定結果(is_tail)を作って、その二つを元にどのパターンかを返答します。

コード

def solve(N, M, S, T):
    is_head = int(S == T[:N])
    is_tail = int(S == T[-N:])
    ans = ((1 - is_head) << 1) | (1 - is_tail)

    return ans


def main():
    N, M = map(int, input().split())
    S = input()
    T = input()
    ans = solve(N, M, S, T)
    print(ans)


if __name__ == "__main__":
    main()

C問題: Festival

atcoder.jp

問題の概要

$N$ 日間のお祭りにおいて、 $A_1, A_2, \cdots, A_M$ 日目に花火が上がる。 $i = 1, 2, \cdots, N$ に対して、 $i$ 日目以降で初めて花火が上がるのは $i$ 日目から数えて何日後か求めよ。

問題の考察

各 $i$ 日目について二分探索することで初めて花火が上がる日を調べます。結果として、計算量は $\mathcal{O} (N \log M)$ となります。

コード

from bisect import bisect_left


def solve(N, M, A):
    ans = []
    for i in range(1, N + 1):
        j = bisect_left(A, i)
        ans.append(A[j] - i)

    return ans


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


if __name__ == "__main__":
    main()

D問題: Polyomino

atcoder.jp

問題の概要

3つのポリオミノをそれぞれ平行移動と90度回転させることで、4x4のグリッドをポリオミノでちょうど敷き詰められるか判定せよ。

問題の考察

この手の問題は集合の論理和と論理積を使うと実装が楽になる印象があります。まずは、ポリオミノがあるマスの集合でもって各ポリオミノを表します。次に、ポリオミノについて平行移動16パターンと回転4パターンを全探索する関数を作ります。最後にその関数を3つのポリオミノにそれぞれ適用し、ポリオミノ同士の論理積が空集合なこと、および、3つのポリオミノの論理和が4x4のグリッドになっていることを判定します。

コード

def simplify(P):
    min_j = min([j for j, _ in P])
    min_k = min([k for _, k in P])
    Q = set([(j - min_j, k - min_k) for j, k in P])

    return Q


def rotate(P):
    Q = set([(-k, j) for j, k in P])

    return Q


def shift(P, dj, dk):
    Q = set([(j + dj, k + dk) for j, k in P])

    return Q


def move(P):
    Q = P
    for _ in range(4):
        Q = simplify(Q)
        for dj in range(4):
            for dk in range(4):
                R = shift(Q, dj, dk)
                yield R
        Q = rotate(Q)


def is_valid(P1, P2, P3):
    if P1 & P2 or P2 & P3 or P1 & P3:
        return False

    TARGET = set([(j, k) for j in range(4) for k in range(4)])
    result = P1 | P2 | P3
    return result == TARGET


def solve(P1, P2, P3):
    for Q1 in move(P1):
        for Q2 in move(P2):
            for Q3 in move(P3):
                if is_valid(Q1, Q2, Q3):
                    return True

    return False


def main():
    Ps = []
    for _ in range(3):
        P = set()
        for j in range(4):
            line = input()
            for k, p in enumerate(line):
                if p == "#":
                    P.add((j, k))
        Ps.append(P)
    ans = "Yes" if solve(*Ps) else "No"
    print(ans)


if __name__ == "__main__":
    main()

2023-09-25

【AtCoder】PythonでABC321を解く

はじめに

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

はじめに
A問題: 321-like Checker
B問題: Cutoff
C問題: 321-like Searcher
D問題: Set Menu

A問題: 321-like Checker

atcoder.jp

問題の概要

正整数の桁が上から順に狭義単調減少になっているとき、その正整数を321-like Numberと呼ぶ。与えられた正整数が321-like Numberかを判定せよ。

問題の考察

itertoolsのpairwiseを使うことで隣り合う2つの要素を取り出せるため、これを使って単調減少しているかを判定します。

コード

from itertools import pairwise


def solve(N):
    return all(s1 > s2 for s1, s2 in pairwise(N))


def main():
    N = input()
    ans = "Yes" if solve(N) else "No"
    print(ans)


if __name__ == "__main__":
    main()

B問題: Cutoff

atcoder.jp

問題の概要

$N$ ラウンドからなる試験の最終結果は、最高スコアのラウンドと最低スコアのラウンドを除いた $N-2$ ラウンド分のスコアの合計となる。 $N-1$ ラウンドまでが終了して各ラウンドのスコアが $A_i$ であるとき、最終結果を $X$ 以上にするために最終ラウンドで取るべきスコアの最小値を求めよ。

問題の考察

最終ラウンドのスコアを $[0, 100$ ]の範囲で全探索します。

コード

def solve(N, X, A):
    ans = float("inf")
    for a in range(101):
        A.append(a)
        total = sum(A) - min(A) - max(A)
        if total >= X:
            ans = min(ans, a)
        A.pop()

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

    return ans


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


if __name__ == "__main__":
    main()

C問題: 321-like Searcher

atcoder.jp

問題の概要

$K$ 番目に小さい321-like Numberを求めよ。

問題の考察

最も大きい321-like Numberは9876543210であること、および、321-like Numberである条件が厳しいことから、321-like Numberを満たす正整数の数は少なそうなため、全てを列挙しても大丈夫そうです。dequeを321-like Numberである1から9で初期化し、①dequeの先頭から正整数を取り出す、②その正整数から作れる321-like Numberをdequeの末尾に入れる、という①②をdequeが空になるまで繰り返します。

コード

from collections import deque


def solve(K):
    fifo = deque(range(1, 10))

    A = []
    while fifo:
        x = fifo.popleft()
        A.append(x)
        for d in range(x % 10):
            xx = 10 * x + d
            fifo.append(xx)

    return A[K]


def main():
    K = int(input())
    ans = solve(K - 1)
    print(ans)


if __name__ == "__main__":
    main()

atcoder.jp

問題の概要

価格 $A_i$ な $N$ 種類の主菜と価格 $B_j$ な $M$ 種類の副菜がある。主菜と副菜からなる定食を提供するとき、その価格は $\min (A_i + B_j, P)$ となる。 $N M$ 通りある定食についての価格の総和を求めよ。

問題の考察

主菜 $A_i$ を全探索するとき、 $B_j < P - A_i$ となる副菜と $B_j \ge P - A_i$ となる副菜の個数が高速に分かればよさそうです。事前に副菜を価格でソートしておき、副菜の価格の累積和を取っておきます。ある主菜 $A_i$ について二分探索で $B_j < P - A_i$ となる副菜の数 $z$ を計算すれば、ある主菜と全副菜との価格の和は累積和を使って $\mathcal O (1)$ で計算できます。副菜のソートに $\mathcal{O} ( M \log M )$ 、各主菜での副菜の数の計算に $\mathcal{O} ( N \log M)$ だけ掛かるため、全体として $\mathcal{O} ( M \log M + N \log M)$ の計算量で価格の総和を求められました。

コード

from bisect import bisect_left
from itertools import accumulate


def solve(N, M, P, A, B):
    B = sorted(B)
    accB = list(accumulate(B, initial=0))

    ans = 0
    for a in A:
        j = bisect_left(B, P - a)
        ans += (j * a + accB[j]) + (M - j) * P

    return ans


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


if __name__ == "__main__":
    main()