【Python】AtCoderの言語アップデートでの変更点をまとめてみたよ。

今年8月にAtCoderで言語アップデートが入りました。(AtCoder 使用できる言語とライブラリの一覧 )
この記事では、AtCoderでのPythonがどのように変わったのかをまとめます。

どのPythonを選べばいいの？

Pythonが4種類おいてありますが、コードテストでいくつかコードを書いて実験してみると、たぶんPython(PyPy3.10-v7.3.12)が一番速そうでした。
基本的にはPython(PyPy3.10-v7.3.12)を選択して提出すればいいのかなと思います。

ac-library-python

元となっているGitHubのリンク → ac-library-python
よくつかうUnionFindやセグメントツリーなどが、手元にライブラリを用意しなくてもつかえるようになりました。
GitHubのリンクにインストール方法からつかい方や中身のコードまで全部書いてありますが、ここでも簡単に説明していきます。

インストール方法

これをコンソールにコピペしてEnterキーを押すだけです。

pip install git+https://github.com/not522/ac-library-python

使用例1 UnionFind

この問題を解いてみます。

$N$ 頂点 $0$ 辺の無向グラフに $Q$ 個のクエリが飛んできます。処理してください。

• $0 \quad u \quad v$ : 辺 $(u,v)$ を追加する。
• $1 \quad u \quad v$ : $(u,v)$ が連結ならば $1$ 、そうでないなら $0$ を出力する。

from atcoder.dsu import DSU

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

for _ in range(Q):
    t, u, v = map(int, input().split())
    if t == 0:
        uf.merge(u, v)
    else:
        if uf.same(u, v):
            print("1")
        else:
            print("0")

使用例2 セグメントツリー

この問題を解いてみます。

長さ $N$ の整数列 $A=(A_1, A_2, \cdots ,A_N)$ があります。
この数列に対して、 $Q$ 個のクエリを処理してください。 $i$ 番目のクエリでは、まず整数 $T_i$ が与えられます。その後、 $T_i$ に応じて以下の操作を行ってください。

• $T_i=1$ : 整数 $X_i,V_i$ が与えられる。 $A_{X_i}$ を $V_i$ で置き換える。
• $T_i = 2$ : 整数 $L_i, R_i$ が与えられる。 $A_{L_i}, A_{L_i+1}, \cdots ,A_{R_i}$ の中での最大値を求める。
• $T_i=3$ : 整数 $X_i,V_i$ が与えられる。 $X_i \leq j \leq N, V_i \leq A_j$ を満たす最小の $j$ を求める。そのような $j$ が存在しない場合、 $j=N+1$ と答える。

from atcoder.segtree import SegTree

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

# SegTree(関数, 単位元, 元となるリスト)
st = SegTree(max, -1, A)

for _ in range(Q):
    t, *q = map(int, input().split())
    if t == 1:
        x, v = q
        x -= 1
        st.set(x, v)
    elif t == 2:
        l, r = q
        ans = st.prod(l - 1, r)
        print(ans)
    elif t == 3:
        x, v = q
        x -= 1
        ans = st.max_right(x, lambda el: el < v) + 1
        print(ans)

math.gcd(), math.lcm()

最大公約数を求める関数 math.gcd() が、引数3つ以上でもつかえるようになりました。

import math

x = math.gcd(30, 60, 105)
print(x)  # 15

最小公倍数を求める関数 math.lcm() が新たに追加されました。

import math

y = math.lcm(24, 48, 120)
print(y)  # 240

# 今まではmath.gcd()をつかって算出していました
# GCD = math.gcd(30, 45)
# LCM = 30 * 45 // GCD

パターンマッチ構文

パターンマッチ構文がつかえるようになりました。
3種類くらいあるクエリがたくさん飛んでくるときなんかにつかえます。
こんな感じのコードになります。

Q = int(input())
for _ in range(Q):
    t, *q = map(int, input().split())
    match t:
        case 1:
            # t=1のときの処理
            pass
        case 2:
            # t=2のときの処理
            pass
        case _:
            # 上の条件に1つも当てはまらなかった(t!=1 and t!=2)ときの処理
            pass

タプルやデータ型でもパターンマッチできます(競プロではあまりつかわなさそうなので省略します)。

sortedcontainers

今まで多重集合(C++だとstd::set)はPythonには標準装備されていなくて、各々で自作したものをつかったり、tatyamさんのSortedSet(Qiita - Python で std::set の代替物を作った)を拝借したりしてました。
今回のアップデートで、sortedcontainersにあるSortedSetがつかえるようになり、わざわざコピペしてくる必要がなくなりました。
下のコードは、ABC217-D Cutting Woodsのコード例です。

from sortedcontainers import SortedSet

l, q = map(int,input().split())
ss = SortedSet()
ss.add(0)
ss.add(l)
for _ in range(q):
    c, x = map(int,input().split())
    if c == 1:
        ss.add(x)
    else:
        idx = ss.bisect_left(x)
        print(ss[idx] - ss[idx - 1])

tatyamさんのSortedSetにあったgt(x)やle(x)のメソッドはなくなっていて、代わりにbisect_leftなどでインデックスを求めないといけなさそうです。
その他のadd, discard, count, indexのメソッドはsortedcontainersのSortedSetでもつかえます。
また、tatyamさんのSortedMultiSet(SortedSetの重複ありバージョン)は、sortedcontainersのSortedListに対応しています。

f-string

実は今までf-stringはコメントアウトしていてもPyPyだとREが出てしまっていました。
今回のアップデートでf-stringが導入されたので、普通につかうことができるようになりました。
個人的に、「このコード、どこでバグが出てるんだろう...」となったときによくつかっています。

x = 101
y = 333
print(f'x+yを計算すると、{x+y}になります')  # x+yを計算すると、434になります
print(f'{x+y=}')  # x+y=434

int.bit_count()

ある整数を2進数表記したときに、そこに含まれている1の数を返すメソッドです。

x, y, z = 13, 30, 255
print(x.bit_count())  # 3 (13の2進数表記は1101)
print(y.bit_count())  # 4 (30の2進数表記は11110)
print(z.bit_count())  # 8 (255の2進数表記は11111111)

実はアップデートの前から存在していたメソッドですが、このメソッドの処理速度がはやくなりました。
たとえばこの問題(ABC258-G - Triangle)は、2進数表記をしたときに含まれている1の数を高速に求めるのがキーポイントです(bitset高速化と言われます)。C++では標準装備されているbitsetを用いて求められますが、Pythonにはそれがなく、63bit以下の整数を2進数表記したときに含まれている1の数を高速に求めるための関数を手元に持っておく必要がありました(以前のbit_countメソッドでは実行時間に間に合いませんでした)。

# 63bit以下の整数を2進数表記したときに立っているbit数
def popcount64(x):
    x = x - ((x >> 1) & 0x5555555555555555)
    x = (x & 0x3333333333333333) + ((x >> 2) & 0x3333333333333333)
    x = (x + (x >> 4)) & 0x0f0f0f0f0f0f0f0f
    x = x + (x >> 8)
    x = x + (x >> 16)
    x = x + (x >> 32)
    return x & 0x7f

今回のアップデートからbit_count()メソッドが最適化されたので、bit_count()をつかうのがラクそうです。

bitarray (おまけ1)

今回のアップデートでbitsetっぽいものとしてbitarrayが追加されました。

from bitarray import bitarray

x = bitarray('1001011')
print(x[2])  # 0
y = bitarray('0100001')
print(x ^ y)  # bitarray('1101010')  (x xor yの計算)
print(y.count(1))  # 2

最後の行は1の数を数えてくれるメソッドです。これが高速ならC++のbitsetとおなじようにつかえたんですが、コードテストで実験してみたところ、これたぶん遅いです。bitset高速化の問題ではつかえなさそうです。

more-itertools (おまけ2)

名前の通り、itertoolモジュールにはないイテレータがたくさんあります。
たとえばこんなのがあります。

from more_itertools import chunked

# リストを3個区切りで分割したいなってときにつかうchunked
chunk = chunked([3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8], 3)
for lst in chunk:
    print(lst)

# 出力結果
# [3, 1, 4]
# [1, 5, 9]
# [2, 6, 5]
# [3, 5, 8]

from more_itertools import distribute

# 4つとばしで要素を見れるdistribute
d1, d2, d3, d4 = distribute(4, [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8])
print(list(d1))  # [3, 5, 5]
print(list(d2))  # [1, 9, 3]
print(list(d3))  # [4, 2, 5]
print(list(d4))  # [1, 6, 8]

他にもたくさんあるんですが、これらを知ってないと解けない問題もあんまりなさそうなので、興味があったらここ(公式APIリファレンス)を眺めてみるくらいでいいと思います。