AtCoderの最新の言語アップデートが過去問にも適用された。この言語アップデートでSageMathを希望していて、採用されている。これで誰も使わなかったら、申し訳ないし、次の言語アップデートで消されてしまうだろう。ということで、紹介記事を書く。でも、私も詳しいわけではないので、皆使って知見を共有してほしい。正解者数2桁くらいの難しい問題で、「え？ SageMathのこの関数を使うだけでしたが、何か？」と誰かが通している姿を見たい。
SageMathとは？
Wikipediaを見ると良いだろうか。要はOSS版Mathematicaである。
文法はほぼPython。競技プログラミング視点で見ると、便利なライブラリが大量に入ったPythonとして使える。
ローカルで動かす方法
AtCoderのコードテストだけだと面倒。
ローカルで動かすには、Dockerを使うのが簡単。
ソースコードをA.sageとして、Windowsなら、
docker run --rm -it -v %CD%:/host sagemath:sagemath sage /host/A.sage
Linuxなら、
docker run --rm -it -v $PWD:/host sagemath:sagemath sage /host/A.sage
色々試すときには、インタラクティブシェルも使える。
docker run --rm -it sagemath/sagemath
これはSageMathに限った話ではないけれど、インタラクティブシェルでは dir(xxx) や help(xxx) が便利。
デフォルトの配色は数字の色が濃い青で見づらいので、
%colors Linux
と打ち込んで配色を変えると良い。
昔はなんかインストールが面倒だった覚えがあって「Dockerが簡単」と言っていたけど、でも、言語アップデートで sudo apt update && sudo apt install -y sagemath で入っているな。普通にインストールしたほうが楽かもしれない。
Pythonとの違い
ほぼPythonなのでPythonのコードがほぼそのまま使える。引っ掛かりそうなところを挙げる。他にもあるかもしれない。
^ が累乗
sage: 2^2
4
数式の記法に合わせているのだろう。** も変わらず使えるから、累乗の計算をしたいときには困らない。排他的論理和を計算したいときに困る。排他的論理和は ^^ になっている。
sage: 2^^2
0
数値リテラルが int ではない
sage: type(123)
<class 'sage.rings.integer.Integer'>
例えば (123).bit_length() がエラーになる……と思っていたけどならないな。昔エラーになった覚えがあるが……。むしろ 123.bit_length() と書けるらしい。Pythonより便利かもしれない。
sage: 123.bit_length()
7
まあ、 int でなくて何か困ることがあったら、 int(123) とキャストすれば良い。
謎の文法 hoge.<fuga>
sage: hoge.<fuga> = PolynomialRing(QQ)
sage: hoge
Univariate Polynomial Ring in fuga over Rational Field
sage: fuga
fuga
名前が分からない（そもそも名前があるのかも分からない）ので、ググれない。まあ、サンプルを見ながら使う分には困らないだろう。
デメリット
先にデメリットを書いておく。
遅い
例えば、この前のABCのA問題の、ほとんど print だけのようなコードで実行時間が1秒。
M = {
"tourist": 3858,
"ksun48": 3679,
"Benq": 3658,
"Um_nik": 3648,
"apiad": 3638,
"Stonefeang": 3630,
"ecnerwala": 3613,
"mnbvmar": 3555,
"newbiedmy": 3516,
"semiexp": 3481,
}
print(M[input()])
ついでにメモリ使用量は約200MB。
これが一番のネック。たいていの問題の制限時間の2秒には収まっているし、たぶん起動が遅いだけだとは思うが、それでも使える時間が半分になってしまうのが厳しい。
起動時間を除いても、素のPythonのように「PyPyで提出したら速くなるかも」ができないのが厳しい。
各種ライブラリは、当然最適化はされているだろうけど、それは10分を9分にするところを頑張っていて、1秒を0.9秒にするところではないのではという気がする。
エディタなどのサポートが無いかも
VSCodeでシンタックスハイライトがされないし、拡張も見当たらない。Qiitaも全部白文字。地味に不便。AtCoderのエディタはシンタックスハイライトが効いた。すごい。
文法がほぼPythonなのだから、拡張子を.pyにすれば良いと思いますよね。SageMathに拡張子が.pyのファイルを渡すと、Pythonとして解釈されてしまう。
機能が多すぎて探せない
sage: len(dir())
1899
グローバルに2,000個近くの変数や関数があった。ドキュメントも、カテゴリ別の1カテゴリだけで1,000ページ近くあったりする。
短時間のコンテスト中に、「この難しいアルゴリズムはSageMathにあるかも！」となったところで、探して使い方を把握できるのかという……。
例
実際にいくつか問題を解いてみる。
有限体
良く見る「mod 998244353で求めてください」というやつ。g = GF(998244353) として有限体を作り、各要素は x = g(123) のように取り出す。あとは x で普通に計算をすれば良い。
sage: g = GF(7)
sage: g(10)
3
sage: g(3)+5
1
sage: g(4)/3
6
sage: g(6)*3
4
sage: g(3)^999999999
6
sage: g(3)**999999999
6
累乗も高速。
実際の問題だとこんな感じ。
N = int(input())
S = input()
g = GF(10^9+7)
T = [g(1)]+[g(0)]*7
for s in S:
P = T
T = [g(0)]*8
for i in range(8):
T[i] += P[i]
if i<7 and s=="atcoder"[i]:
T[i+1] += P[i]
print(T[7])
素因数分解
これは得意。ちゃんと高速なアルゴリズムが実装されているので、128 bitくらいなら一瞬で素因数分解できる。
sage: factor(171199332458954054723586764520404288963)
10221620570433624431 * 16748746569027790573
sage: x = factor(12)
sage: x
2^2 * 3
sage: x[0]
(2, 2)
sage: x[1]
(3, 1)
素因数分解を試し割りでするなら一工夫必要な問題。SageMathなら素直に素因数分解できる。素因数分解が難しいことが現代の暗号の基礎に～という話があるけれど、AtCoderで出てくる整数は基本的に64 bitに収まるので、AtCoderでは素因数分解は簡単な問題である。
T = int(input())
for _ in range(T):
N = int(input())
pq = factor(N)
if pq[0][1]==2:
print(pq[0][0], pq[1][0])
else:
print(pq[1][0], pq[0][0])
いや、でも、最初のテストケースで2.5秒掛かっているな……。制限時間が2秒だとTLE。
https://atcoder.jp/contests/abc284/submissions/45513585
もう1回投稿したら0.6秒になった。
https://atcoder.jp/contests/abc284/submissions/45513599
テストを実行するVMが温まったら速くなるとかあるのだろうか。この不安定さはコンテスト中に使うのが厳しい……。
行列
matrix([[1, 2], [3, 4]]) として行列が作れる。 matrix(GF(10^9+7), [[1, 2], [3, 4]]) なら mod 1000000007 で計算。
行列累乗で（も）解ける問題。
S = input()
K = int(input())
# https://atcoder.jp/contests/abc200/editorial/1249
M = 10^9+7
zero = matrix(GF(M), [
[1, 0, 0, 0, 0],
[1, 0, 0, 0, 0],
[0, 0, 0, 1, 1],
[0, 0, 0, 1, 0],
[0, 0, 0, 0, 1],
])
one = matrix(GF(M), [
[0, 1, 0, 0, 1],
[0, 1, 0, 0, 0],
[0, 0, 1, 0, 0],
[0, 0, 1, 0, 0],
[0, 0, 0, 0, 1],
])
c2m = {
"0": zero,
"1": one,
"?": zero+one,
}
m = matrix.identity(GF(M), 5)
for s in S[1:]:
m *= c2m[s]
m = m*(c2m[S[0]]*m)^(K-1)
if S[0]=="0":
print(m[0][4])
if S[0]=="1":
print(m[2][4])
if S[0]=="?":
print(m[0][4]+m[2][4])
TLE。残念。
文字列
Word("abc") に文字列関連のアルゴリズムがある。
最長の回文を探す問題。想定解法は単純な $O\left(\left|S\right|^3\right)$ のアルゴリズムだが、回文については色々とアルゴリズムがあり、SageMathにも実装されている。
print(max(Word(input()).lps_lengths()))
これは余裕で最短コードだろと思ったら、上にNibblesという言語がいた。コードゴルフ用の言語らしい。なるほど……。
グラフ
解説で名前が出てくるようなアルゴリズムはだいたいありそう。
最短経路。
N, M = map(int, input().split())
G = DiGraph()
for _ in range(M):
u, v, c = map(int, input().split())
G.add_edge(u, v, {"weight": c})
print(G.shortest_path_length(0, N-1, weight_function=lambda e: e[2]["weight"]))
多項式
多項式の除算をそのまま書けば良い。ZZ は整数（$\mathbb{Z}$）上の多項式だということ。
N, M = map(int, input().split())
A = list(map(int, input().split()))
C = list(map(int, input().split()))
R.<x> = PolynomialRing(ZZ)
a = sum(A[i]*x^i for i in range(N+1))
c = sum(C[i]*x^i for i in range(N+M+1))
b = c//a
B = [b[i] for i in range(M+1)]
print(*B)
ラグランジュ補間もあるので、この問題がさくっと解けそう。
……と思ったけれど、TLE。
N, M, K = map(int, input().split())
# https://atcoder.jp/contests/abc208/editorial/2195
MOD = 10^9+7
F = [pow(n, K, MOD) for n in range(K+M+1)]
for _ in range(M):
for i in range(1, K+M+1):
F[i] = (F[i-1]+F[i])%MOD
R.<x> = PolynomialRing(GF(MOD))
f = R.lagrange_polynomial([(n, F[n]) for n in range(K+M+1)])
print(f(N))
数独
ドキュメントを眺めていたら、数独があった。もしAtCoderで数独の問題が出たら無双できる。
sage: a = Sudoku('5...8..49...5...3..673....115..........2.8..........187....415..3...2...49..5...3')
print(a)
+-----+-----+-----+
|5 | 8 | 4 9|
| |5 | 3 |
| 6 7|3 | 1|
+-----+-----+-----+
|1 5 | | |
| |2 8| |
| | | 1 8|
+-----+-----+-----+
|7 | 4|1 5 |
| 3 | 2| |
|4 9 | 5 | 3|
+-----+-----+-----+
sage: print(next(a.solve()))
+-----+-----+-----+
|5 1 3|6 8 7|2 4 9|
|8 4 9|5 2 1|6 3 7|
|2 6 7|3 4 9|5 8 1|
+-----+-----+-----+
|1 5 8|4 6 3|9 7 2|
|9 7 4|2 1 8|3 6 5|
|3 2 6|7 9 5|4 1 8|
+-----+-----+-----+
|7 8 2|9 3 4|1 5 6|
|6 3 5|1 7 2|8 9 4|
|4 9 1|8 5 6|7 2 3|
+-----+-----+-----+