ABC312をPythonで解いてみたよ。(A～F問題)

AtCoder Beginners Contest 312 (ABC312) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。

A - Chord

問題ページ

考察

テクニカルな解法はTwitterにも公式の解説にもあるとは思いますが、気合いで解けるならそれで十分okです。
「この中に文字列 $S$ は存在しますか？」の問題はリストよりもセットの方が高速ですが、たかが7個の判定なので正直どっちでもいいです。

コード

lst = [
    "ACE",
    "BDF",
    "CEG",
    "DFA",
    "EGB",
    "FAC",
    "GBD"
]
S = input()
if S in lst:
    print("Yes")
else:
    print("No")

B - TaK Code

問題ページ

考察

「 $9 \times 9$ のマスが与えられます。左上 $3 \times 3$ マスと右下 $3 \times 3$ マスが'#'で、その周囲が'.'で囲まれているならTrueを、そうでないならFalseを返してください。」を解く関数をつくります(下のコードのis_tak関数)
この関数さえつくれたら、あとはすべてのマスについてis_tak関数をつかって条件を満たしているかどうかを判定すればokです。

コード

N, M = map(int, input().split())
S = [input() for _ in range(N)]


def is_tak(si, sj):
    for i in range(si, si + 3):
        for j in range(sj, sj + 3):
            if S[i][j] != '#':
                return False
    for i in range(si + 6, si + 9):
        for j in range(sj + 6, sj + 9):
            if S[i][j] != '#':
                return False
    for i in range(si, si + 4):
        if S[i][sj + 3] != '.':
            return False
    for i in range(si + 5, si + 9):
        if S[i][sj + 5] != '.':
            return False
    for j in range(sj, sj + 4):
        if S[si + 3][j] != '.':
            return False
    for j in range(sj + 5, sj + 9):
        if S[si + 5][j] != '.':
            return False
    return True


for i in range(N - 8):
    for j in range(M - 8):
        if is_tak(i, j):
            print(i + 1, j + 1)

C - Invisible Hand

問題ページ

考察

リスト $A$ の中で値が $X$ 以下の要素の個数を$C_A$ 、リスト $B$ の中で値が $X$ 以上の要素の個数を $C_B$ としたとき、$C_A \geqq C_B$ ならokで、これは値 $X$ が小さくなるほど条件を満たすのがしんどくなる性質があって、どこまで $X$ を減らせますか？という問題です。
こういうタイプの、「 $X$ が大きければTrueで、そこから徐々に $X$ を減らしていくとどこかでFalseになって、そこからはいくら減らしてもずっとFalse」みたいな構造をしている問題には、二分探索が有効です。

下のコードでは、 $ok=10^9+1, ng=0$ として解いています(めぐる式二分探索)。

コード

from bisect import bisect_right, bisect_left

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


def judge(m):
    c1 = bisect_right(A, m)
    c2 = M - bisect_left(B, m)
    return c1 >= c2


ok, ng = 10 ** 9 + 1, 0
while ok - ng > 1:
    mid = (ok + ng) // 2
    if judge(mid):
        ok = mid
    else:
        ng = mid
print(ok)

D - Count Bracket Sequences

問題ページ

考察

文字列 $S$ の長さが奇数のとき、 '(' の数と ')' の数が一緒にならないので、どうやっても条件を満たせません。先に下のコードを書いて長さが奇数のものを除外しておきます。

先に例外やコーナーケースを書くことへのおきもち

(これを書かなくてもどうせACできたじゃん！みたいな問題も多数ありますが、書いてもACできるし考察がすっきりするので、奇数のを除外したり何かの値が $0$ のものを除外したりするのはたとえ意味がなかったとしても先にやっといていいと思います。)

if len(S) % 2 == 1:
    print(0)
    exit()

正しい括弧列についての条件が問題文中にありますが、その条件は下の $2$ つの条件に言い換えられます。

• '(' と')' の数が等しい。
• 左から順番に見ていって、今まで見た文字の中で '('の数は ')'の数と同じかそれ以上。

これは過去問で何度も出ている典型なので、おぼえておくのがオススメです！
$2$ つめの条件の通りに、左からDPで解いていきます。
具体的には、次のリストを用意します。

• $dp[i][j]$ : 左から $i$ 個目の文字までで、'(' の数が $j$ 個になるパターンの数
• $dp[0][0]=1$

そして、 $2$ つめの条件に気をつけながら遷移していきます。(下のコードのis_ket関数に書いてあります)

$1$ つめの条件より、答えは $dp[len(S)][len(S)//2]$ になります。

コード

MOD = 998244353
S = input()

if len(S) % 2 == 1:
    print(0)
    exit()

# dp[i][j]: i個目まで見て、'('がj個のパターン数
dp = [[0] * 3001 for _ in range(len(S) + 1)]
dp[0][0] = 1


def is_bra(i):
    for j in range(3000):
        dp[i + 1][j + 1] += dp[i][j]
        if dp[i + 1][j + 1] >= MOD:
            dp[i + 1][j + 1] -= MOD


def is_ket(i):
    for j in range(3001):
        prev_bra = j
        now_ket = i + 1 - prev_bra
        if prev_bra < now_ket:
            continue
        dp[i + 1][j] += dp[i][j]
        if dp[i + 1][j] >= MOD:
            dp[i + 1][j] -= MOD


for i, s in enumerate(S):
    if s == '(':
        is_bra(i)
    elif s == ')':
        is_ket(i)
    elif s == '?':
        is_bra(i)
        is_ket(i)

print(dp[len(S)][len(S) // 2])

E - Tangency of Cuboids

問題ページ

考察

$2$ 次元や $3$ 次元の問題になったときの典型として、「次元を $1$ つ減らした問題ならどう解きますか？」の問題を考えてみるのがあります。
問題で与えられた数字のままだとそれぞれの直方体がどう接するかもよく分からないので、いったん次の問題に言い換えてみます。

それぞれの直方体について、「xy平面と平行な向き」で接する直方体の数を求めてください。

これがもし解けたら、yz平面・zx平面でも同じように解けばいいので、ACできそうです。
愚直にやるとしたら、次のような解法があります。

• $101^3$ のサイズの $3$ 次元リスト$Grid$を用意する。任意の$x,y,z$に対して$Grid[x][y][z]$はset()になっている。
• 直方体それぞれについて、接している面すべてで$Grid[x][y][z]$に直方体の番号を格納する。
• $Grid$ 内のすべてのsetについて、もし要素数が$2$以上なら、ペアを記録するためのセットpair_setにありうるペアをすべて格納する。

これの問題点は、$Grid$ 内のsetについて、たとえば要素数が $N$ 個だったらペアの個数は $N \times (N-1) / 2$ でやばいという点です。
しかし、Grid内のsetについて、その要素数は最大でも $2$ つです。なぜなら、問題文中にこう書いてあるからです。

直方体同士は重なっていません。厳密には、相異なるどの $2$ つの直方体の共通部分の体積も $0$ です。

直方体の上面と、別の直方体の下面が接していたとして、さらに別の直方体が同じ面で接するようにしようとすると、どうしてもどちらかの直方体と重なってしまいます(たぶん実際にサイコロとかを用意して試してみた方が分かりやすいです)。
なので、さっきの愚直っぽい解法で解けます。接した面それぞれでsetを持つ必要はないので、次のように変更しておきます。

• $101^3$ のサイズの $3$ 次元リスト$Grid$を用意する。任意の$x,y,z$に対して$Grid[x][y][z]$は $-1$になっている。
• 直方体それぞれについて、接している面すべてについて、もし$Grid[x][y][z]=-1$ならば、$Grid[x][y][z]$に直方体の番号を代入する。
• もし $Grid[x][y][z] \neq -1$ ならば、それは別の直方体が接している面だということになるので、ペアを記録するためのセットpair_setにペアを格納する。

コード

N = int(input())
cube = [list(map(int, input().split())) for _ in range(N)]

lst1 = []
lst2 = []
lst3 = []
for x1, y1, z1, x2, y2, z2 in cube:
    lst1.append((x1, x2, y1, y2, z1, z2))
    lst2.append((x1, x2, z1, z2, y1, y2))
    lst3.append((y1, y2, z1, z2, x1, x2))

pair_set = set()


def solve(lst):
    Grid = [[[-1] * 101 for _ in range(100)] for _ in range(100)]
    for i in range(N):
        x1, x2, y1, y2, z1, z2 = lst[i]
        for x in range(x1, x2):
            for y in range(y1, y2):
                for z in (z1, z2):
                    if Grid[x][y][z] == -1:
                        Grid[x][y][z] = i
                    else:
                        if Grid[x][y][z] < i:
                            pair_set.add((Grid[x][y][z], i))
                        else:
                            pair_set.add((i, Grid[x][y][z]))


for lst in [lst1, lst2, lst3]:
    solve(lst)

ans_list = [0] * N
for a, b in pair_set:
    ans_list[a] += 1
    ans_list[b] += 1
for ans in ans_list:
    print(ans)

F - Cans and Openers

問題ページ

考察

缶切りをつかわなかったら、 $T=0$ の商品を優先度の大きいものから順に $M$ 個選べばよいです。
その $M$ 個が手元にあるとして、そこに缶切りを $1$ つだけ追加するとします。とりあえず缶切りを $1$ つつかうので、手元にあった $T=0$ の商品 $M$ 個の中で優先度の最も低いものを $1$ つ捨てます。手元には $M-1$ 個の商品がある状態です。
その缶切りが例えば $5$ 個までの商品につかえるとしたら、$T=1$ の商品を優先度の大きいものから順に $5$ つ選んで、手元の商品に追加していきます。ただし、手元の商品の個数は $M-1$ 個まで(缶切りの分で $1$ 減ってる)なので、$1$ つ追加するごとに 優先度の最も低い商品 $1$ つを捨てなければなりません。
これは、下の $2$ つができる構造があるとめちゃくちゃうれしくなります。

• 要素の追加がはやくできる
• リスト内の最小値をすぐに取り出せる

これは優先度付きキューでクリアできます。

$T=0,1,2$の商品をそれぞれ別のリストに格納して、それぞれのリストをソートしておきます(缶切りは たくさんの缶を開けられるものからつかいたいし、 $T=0,1$ の商品も優先度の高い商品から順番に見ていきたいので)。

手元の商品を入れるヒープキューを用意します。
最初は缶切りを $1$ つもつかわない場合を考えたくて、 $T=0$ の商品を優先度の大きいものから順に $M$ 個だけ取り出してヒープキューに入れます。
次は缶切りを $1$ つだけつかう場合を考えます。ヒープキューから要素を $1$ つだけ取り出します(缶切りの分)。
そして、その缶切りで開けられる缶の数だけ、$T=1$ の商品を優先度の大きいものから順にヒープキューに入れます。入れた分だけ、ヒープキューから最小値を取り出します。
これを缶切りがなくなるか、 $T=1$ の商品がなくなるまで続けます。
ヒープキュー内の要素の総和を表す変数をなにか用意して適宜いじってあげることで、つかう缶切りの数ごとに答えを更新し続けることでACできます。

下のコードのクラスについて

これは事前にライブラリ化してあるものではなく、本番中に書いてるものです。 今回のB問題やD問題で関数化して実装を分けることで実装・考察が楽になるのと同じで、こうすると個人的に楽かなと思って書いてます。 「これくらいならクラスなんかに分けなくても無理なく書ける！」となる人は無理してクラスを書かなくてもいいし、逆に「コードが多いのは得意じゃない！」という私と同じタイプの人は、コンテスト中でも関数を作る感覚で専用のクラスをつくっちゃってもいいのかなと思います。

コード

from heapq import heappop, heappush


class MyQue:
    def __init__(self, n_max):
        self.que = []
        self.sum = 0
        self.n_max = n_max

    def add(self, x):
        heappush(self.que, x)
        self.sum += x
        if len(self.que) > self.n_max:
            self.discard()

    def discard(self):
        if self.que:
            el = heappop(self.que)
            self.sum -= el

    def n_max_dec(self):
        self.n_max -= 1
        if len(self.que) > self.n_max:
            self.discard()

    def get_sum(self):
        return self.sum


N, M = map(int, input().split())
lst1 = []
lst2 = []
lst3 = []
for _ in range(N):
    t, x = map(int, input().split())
    if t == 0:
        lst1.append(x)
    elif t == 1:
        lst2.append(x)
    else:
        lst3.append(x)
lst1.sort()
lst2.sort()
lst3.sort()

que = MyQue(M)
for _ in range(M):
    if lst1:
        el = lst1.pop()
        que.add(el)
    else:
        break
ans = que.get_sum()

while lst2 and lst3:
    que.n_max_dec()
    el3 = lst3.pop()
    for _ in range(el3):
        if not lst2:
            break
        el2 = lst2.pop()
        que.add(el2)
    ans = max(ans, que.get_sum())

print(ans)