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

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

A - Counting Passes

問題

$L$ 点以上を取った人は何人いますか？

考察

for文をつかって リスト $A$ の要素を1つずつ見て、要素が $L$ 以上だったら変数 $ans$ を $1$ だけプラスします。

コード

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

ans = 0
for a in A:
    if a >= L:
        ans += 1
print(ans)

内包表記をつかった解法

内包表記をつかうと、for文で4行くらいつかっていたコードを1行でスッキリまとめられます。

ここでは内包表記については詳しく書きませんが、「リスト内包表記」などと調べると詳しい記事がたくさん出てきます。

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

ans = len([a for a in A if a >= L])
print(ans)

B - Minimize Abs 1

問題

考察

問題の読解が難しいですが、要するに、リスト $A$ の各要素 $a$ について下の場合分けをすればいいです。

• $a<=L$ のとき、$L$
• $L<a<R$ のとき、 $a$
• $R<=a$ のとき、 $R$

下のコードでは、上の場合分けを関数化して実装しています。

コード

def f(a):
    if a <= L:
        return L
    if R <= a:
        return R
    return a


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

ans_lst = []
for a in A:
    ans_lst.append(f(a))
print(*ans_lst)

C - Minimize Abs 2

問題

$|x^2+y^2-D|$ を最小にしてね。

考察

$x$ が決まっていたとします。
$|x^2+y^2-D|$ を最小にしたいので、$y=\sqrt{|D-x^2|}$ にすればいいです。ただし $y$ は整数なのできちんと$y=\sqrt{|D-x^2|}$ にはならないかもしれないですが、それでもだいたい $y=\sqrt{|D-x^2|}$ になります。
$(2\times 10^6)^2=4\times 10^{12} > 2 \times 10^{12} >=D$ なので、 $x$ を $0$ から $2\times 10^6$ までfor文で動かして、先ほどの計算をすればいいです。

下のコードでは、math.isqrt(abs(D-x*x)) で　$y=\sqrt{|D-x^2|}$ の整数部分を求めて、その $±10$ で $|x^2+y^2-D|$ の値を調べています。
(本来は、その整数部分を $i$ としたとき $i$ と $i+1$ だけ調べればいいですが、コンテストではACするのが優先なので雑に広めにチェックしています。)

コード

import math

D = int(input())

ans = 1 << 60
for x in range(0, 2_000_000):
    yy = math.isqrt(abs(D - x * x))
    for y in range(yy - 10, yy + 11):
        ans = min(ans, abs(x * x + y * y - D))

D - Counting Ls

問題

$3$ つの'o'で、「そのうち $2$ つだけが同じ行にある」と 「そのうち $2$ つだけが同じ列にある」 の条件をどちらも満たしているような組はいくつありますか？
https://atcoder.jp/contests/abc330/tasks/abc330_d

考察

図で表すと、こんな感じです。

こういうL字になるような $3$ つの 'o'を選びたくて、その選び方は何パターンありますか？という問題。

このL字のかどっこで全探索する方針でいきます。
$(i,j)$ の座標が 'o' のとき、

• 行 $i$ のなかで、(i,j) 以外の'o' の個数 ($=r$ とします。)
• 列 $j$ のなかで、(i,j) 以外の'o' の個数 ($=c$ とします。)
  の $2$ つが分かると、 $(i,j)$ がL字のかどっこになるような選び方のパターンは $r \times c$ になります。

各行・各列の'o'の個数をあらかじめ計算しておけば、上の $r \times c$ は$O(1)$ で求められるので、全体で $O(N^2)$ で解くことができます。

コード

N = int(input())
S = [input() for _ in range(N)]

row = [0 for _ in range(N)]
col = [0 for _ in range(N)]
for i in range(N):
    for j in range(N):
        if S[i][j] == "o":
            row[i] += 1
            col[j] += 1

ans = 0
for i in range(N):
    for j in range(N):
        if S[i][j] == "o":
            r = row[i] - 1
            c = col[j] - 1
            ans += r * c
print(ans)

E - Mex and Update

問題

リスト $A$ の一点更新クエリが $Q$ 個与えられます。
クエリが来るたびに、 $A$ のmexを求めてね。

考察

mexは、$0$ 以上の整数の中でリストに含まれないような最小の整数です。
リスト $A$ の長さは $N$ で与えられています。
たとえば $A=[0,1,\cdots N-1]$ のとき、mexは $N$ になります。ちなみに mexの最大値は $N$ になります。

collections.Counter をつかって、リスト $A$ にある要素の数を管理します。
また、SortedSetをつかって、常にリスト $A$ の中にない整数を管理します。

クエリが来るたびに、SortedSetの中身(=リスト $A$ にない整数セット)を調整して、その中で最小の整数を出力します。

コメントもつけているので、細かい実装は下のコードを見てください。

コード

from collections import Counter

"""
tatyamさん作の、SortedSetです。
使わせていただき、ありがとうございます！
https://github.com/tatyam-prime/SortedSet/blob/main/SortedSet.py

・使い方(個人的まとめ)
s=SortedSet()
s.a: SortedSetの中身を返す。
len(s), x in s, x not in s: リストと同じ要領で使える。
s.add(x): xを追加してTrueを返す。ただしxがすでにs内にある場合、xは追加せずにFalseを返す。
s.discard(x): xを削除してTrueを返す。ただしxがs内にない場合、何もせずにFalseを返す。
s.lt(x): xより小さい最大の要素を返す。もし存在しないなら、Noneを返す。
s.le(x): x　以下の　最大の要素を返す。もし存在しないなら、Noneを返す。
s.gt(x): xより大きい最小の要素を返す。もし存在しないなら、Noneを返す。
s.ge(x): x　以上の　最小の要素を返す。もし存在しないなら、Noneを返す。
s.index(x): xより小さい要素の数を返す。
s.index_right(x): x以下の要素の数を返す。

・使い方URL
https://github.com/tatyam-prime/SortedSet
"""
# (SortedSetのコード省略。 上のURLからコピペしてね。)

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

# se: 0～Nの中でリストAの中にない整数のセット。
se = SortedSet([i for i in range(N + 1)])
for a in A:
    se.discard(a)

counter = Counter(A)

for _ in range(Q):
    i, x = map(int, input().split())
    i -= 1
    prev = A[i]
    A[i] = x
    counter[prev] -= 1
    counter[x] += 1
    if counter[prev] == 0:
        # リストAからprevがなくなったので、mexの値の候補にprevが入る。
        se.add(prev)
    # リストAにxがあるので、seの中にxがあった場合はxを取り出す。
    # (SortedSetの中に無い要素をdiscardしても、エラーはおきません。)
    se.discard(x)
    print(se.ge(0))  # seの中にある0以上の最小の整数を出力する。

おまけ

sortedcontainersの中にも、同じ機能を持つSortedSetがあります。

from sortedcontainers import SortedSet

se = SortedSet()

ですが、このSortedSetよりもtatyamさん作のSortedSetの方がはやく、上の解答コードをsortedcontainersのSortedSetに置き換えて提出するとTLEになります。(提出コード(TLE))
基本的にはtatyamさん作のSortedSetをつかう方がいいのかなと思います。