mean_abs_error’s blog

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

【AtCoder】PythonでABC322を解く

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()