ABC294 回想・解説(Python)

前提

AtCoder Biginner Contest294のPythonの解答コードです。

A - Filter

問題文通りに実装する。ループの中で2で割れるかどうかを判定し、条件に合うものを配列に格納する

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

ans = [a for a in A if a % 2 == 0]
print(*ans)

B - ASCII Art

これは2重だが、ループして判定するのみ。

H, W = map(int, input().split())
A = [list(map(int, input().split())) for _ in range(H)]

ans = []
for i in range(H):
    row = ""
    for j in range(W):
        if A[i][j] == 0:
            row += "."
        else:
            row += chr(A[i][j] + ord("A") - 1)
    ans.append(row)

print(*ans, sep="\n")

C - Merge Sequences

Aの配列をループしながら、Aの要素がBのどこにあるかを二分探索した。

pypyでTLEで焦って、Pythonで実行してみるまでに20分くらい時間を無駄にした。

解説を見たら普通に配列をくっつけて、順位をdictで管理していた。

色んな意味で悲c。

from bisect import bisect_left

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

a_idx = ""
b_idx = ""
cur = 0
pre_idx = 0
for a in A:
    idx = bisect_left(B, a)
    if idx - pre_idx == 0:
        pre_idx = idx
        cur += 1
        a_idx += str(cur) + " "
    else:
        diff = idx - pre_idx
        b_idx += " ".join([str(cur + i) for i in range(1, diff + 1)]) + " "
        cur += diff + 1
        a_idx += str(cur) + " "
        pre_idx = idx
b_idx += " ".join([str(cur + i) for i in range(1, M - pre_idx + 1)])
print(a_idx.strip())
print(b_idx.strip())

D - Bank

setで管理したいが、最小値も欲しいので、sorted::set使うことを検討。

Pythonのsorted::setを作ってくれている先人に大いなる感謝をしつつ貼り付け。

素直に実装するとTLE。よくよく問題文を見ると1の場合の状態は管理する必要がない(3は受付に呼ばれていない人は呼ばれない)。

初めからすべての人が呼ばれていることにして(called)、2で受付に行ってしまった人をcalledから取り除いてあげて、3の際にcalledの中から最小値をとってくるように修正。AC。

# https://github.com/tatyam-prime/SortedSet/blob/main/SortedSet.py
import math
from bisect import bisect_left, bisect_right
from typing import Generic, Iterable, Iterator, TypeVar, Union, List

T = TypeVar("T")


class SortedSet(Generic[T]):
    BUCKET_RATIO = 50
    REBUILD_RATIO = 170

    def _build(self, a=None) -> None:
        """Evenly divide `a` into buckets."""
        if a is None:
            a = list(self)
        size = self.size = len(a)
        bucket_size = int(math.ceil(math.sqrt(size / self.BUCKET_RATIO)))
        self.a = [
            a[size * i // bucket_size : size * (i + 1) // bucket_size]
            for i in range(bucket_size)
        ]

    def __init__(self, a: Iterable[T] = []) -> None:
        """Make a new SortedSet from iterable.
        O(N) if sorted and unique / O(N log N)
        """
        a = list(a)
        if not all(a[i] < a[i + 1] for i in range(len(a) - 1)):
            a = sorted(set(a))
        self._build(a)

    def __iter__(self) -> Iterator[T]:
        for i in self.a:
            for j in i:
                yield j

    def __reversed__(self) -> Iterator[T]:
        for i in reversed(self.a):
            for j in reversed(i):
                yield j

    def __len__(self) -> int:
        return self.size

    def __repr__(self) -> str:
        return "SortedSet" + str(self.a)

    def __str__(self) -> str:
        s = str(list(self))
        return "{" + s[1 : len(s) - 1] + "}"

    def _find_bucket(self, x: T) -> List[T]:
        """Find the bucket which should contain x. self must not be empty."""
        for a in self.a:
            if x <= a[-1]:
                return a
        return a

    def __contains__(self, x: T) -> bool:
        if self.size == 0:
            return False
        a = self._find_bucket(x)
        i = bisect_left(a, x)
        return i != len(a) and a[i] == x

    def add(self, x: T) -> bool:
        """Add an element and return True if added. / O(√N)"""
        if self.size == 0:
            self.a = [[x]]
            self.size = 1
            return True
        a = self._find_bucket(x)
        i = bisect_left(a, x)
        if i != len(a) and a[i] == x:
            return False
        a.insert(i, x)
        self.size += 1
        if len(a) > len(self.a) * self.REBUILD_RATIO:
            self._build()
        return True

    def discard(self, x: T) -> bool:
        """Remove an element and return True if removed. / O(√N)"""
        if self.size == 0:
            return False
        a = self._find_bucket(x)
        i = bisect_left(a, x)
        if i == len(a) or a[i] != x:
            return False
        a.pop(i)
        self.size -= 1
        if len(a) == 0:
            self._build()
        return True

    def lt(self, x: T) -> Union[T, None]:
        """Find the largest element < x, or None if it doesn't exist."""
        for a in reversed(self.a):
            if a[0] < x:
                return a[bisect_left(a, x) - 1]

    def le(self, x: T) -> Union[T, None]:
        """Find the largest element <= x, or None if it doesn't exist."""
        for a in reversed(self.a):
            if a[0] <= x:
                return a[bisect_right(a, x) - 1]

    def gt(self, x: T) -> Union[T, None]:
        """Find the smallest element > x, or None if it doesn't exist."""
        for a in self.a:
            if a[-1] > x:
                return a[bisect_right(a, x)]

    def ge(self, x: T) -> Union[T, None]:
        """Find the smallest element >= x, or None if it doesn't exist."""
        for a in self.a:
            if a[-1] >= x:
                return a[bisect_left(a, x)]

    def __getitem__(self, x: int) -> T:
        """Return the x-th element, or IndexError if it doesn't exist."""
        if x < 0:
            x += self.size
        if x < 0:
            raise IndexError
        for a in self.a:
            if x < len(a):
                return a[x]
            x -= len(a)
        raise IndexError

    def index(self, x: T) -> int:
        """Count the number of elements < x."""
        ans = 0
        for a in self.a:
            if a[-1] >= x:
                return ans + bisect_left(a, x)
            ans += len(a)
        return ans

    def index_right(self, x: T) -> int:
        """Count the number of elements <= x."""
        ans = 0
        for a in self.a:
            if a[-1] > x:
                return ans + bisect_right(a, x)
            ans += len(a)
        return ans


N, Q = map(int, input().split())

called = SortedSet(range(1, N + 1))
ans = []
for _ in range(Q):
    (*query,) = map(int, input().split())
    if query[0] == 1:
        continue
    elif query[0] == 2:
        called.discard(query[1])
    elif query[0] == 3:
        _min_num = called.gt(0)
        ans.append(_min_num)
print(*ans, sep="\n")

E - 2xN Grid

Cで時間を使ってしまったために、方針だけを考えてtime up。

連続する数字の最初と最後のインデックスを管理するようにして、被っている区間を計上するようにした。

入力例1の1行目だと{1: (0,1), 3:(2,3), 2:(4,6)...}のような感じ。2行目は{1:(0,3), 2:(4,4)...}のようになっているから1行目の1は0から1,2行目の1は0から3なので、被っている区間は0から1の2つを計上。これを繰り返すことで求まる。

from collections import defaultdict

L, N1, N2 = map(int, input().split())
row1 = defaultdict(list)
row2 = defaultdict(list)

idx = 1
for _ in range(N1):
    a, b = map(int, input().split())
    row1[a].append((idx, idx + b - 1))
    idx += b

idx = 1
for _ in range(N2):
    a, b = map(int, input().split())
    row2[a].append((idx, idx + b - 1))
    idx += b


ans = 0
for k, v in row1.items():
    if k in row2:
        i, j = 0, 0
        while i < len(v) and j < len(row2[k]):
            a, b = v[i]
            c, d = row2[k][j]
            if b < c:
                i += 1
            elif d < a:
                j += 1
            else:
                ans += min(b, d) - max(a, c) + 1
                if b < d:
                    i += 1
                else:
                    j += 1

print(ans)