AtCoderでSageMathを使う

https://atcoder.jp/contests/language-test-202301

https://www.sagemath.org/

AtCoderの最新の言語アップデートが過去問にも適用された。この言語アップデートでSageMathを希望していて、採用されている。これで誰も使わなかったら、申し訳ないし、次の言語アップデートで消されてしまうだろう。ということで、紹介記事を書く。でも、私も詳しいわけではないので、皆使って知見を共有してほしい。正解者数2桁くらいの難しい問題で、「え？　SageMathのこの関数を使うだけでしたが、何か？」と誰かが通している姿を見たい。

# SageMathとは？

Wikipediaを見ると良いだろうか。要はOSS版Mathematicaである。

https://ja.wikipedia.org/wiki/SageMath

文法はほぼ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
```

https://doc.sagemath.org/html/en/reference/rings_standard/sage/rings/integer.html#sage.rings.integer.Integer.bit_length

まあ、 `int` でなくて何か困ることがあったら、 `int(123)` とキャストすれば良い。

# 謎の文法 `hoge.<fuga>`

```
sage: hoge.<fuga> = PolynomialRing(QQ)
sage: hoge
Univariate Polynomial Ring in fuga over Rational Field
sage: fuga
fuga
```

https://doc.sagemath.org/html/ja/tutorial/tour_polynomial.html

名前が分からない（そもそも名前があるのかも分からない）ので、ググれない。まあ、サンプルを見ながら使う分には困らないだろう。

# デメリット

先にデメリットを書いておく。

## 遅い

例えば、この前のABCのA問題の、ほとんど `print` だけのようなコードで実行時間が1秒。

```abc319_a.sage
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()])
```

https://atcoder.jp/contests/abc319/submissions/45473645

ついでにメモリ使用量は約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にあるかも！」となったところで、探して使い方を把握できるのかという……。


# 例

実際にいくつか問題を解いてみる。

## 有限体

https://doc.sagemath.org/html/en/reference/finite_rings/index.html

良く見る「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
```

累乗も高速。

実際の問題だとこんな感じ。

https://atcoder.jp/contests/typical90/tasks/typical90_h

```typical90_h.sage
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])
```

## 素因数分解

https://doc.sagemath.org/html/en/reference/rings_standard/sage/arith/misc.html#sage.arith.misc.factor

これは得意。ちゃんと高速なアルゴリズムが実装されているので、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)
```

https://atcoder.jp/contests/abc284/tasks/abc284_d

素因数分解を試し割りでするなら一工夫必要な問題。SageMathなら素直に素因数分解できる。素因数分解が難しいことが現代の暗号の基礎に～という話があるけれど、AtCoderで出てくる整数は基本的に64 bitに収まるので、AtCoderでは素因数分解は簡単な問題である。

```abc284_d.sage
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。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/25276/55d54e77-f601-1bd1-3f5a-ab38aac60f2d.png)

[https://atcoder.jp/contests/abc284/submissions/45513585](https://atcoder.jp/contests/abc284/submissions/45513585)

もう1回投稿したら0.6秒になった。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/25276/64427564-77c4-8b62-a818-d9989277b6f6.png)

[https://atcoder.jp/contests/abc284/submissions/45513599](https://atcoder.jp/contests/abc284/submissions/45513599)

テストを実行するVMが温まったら速くなるとかあるのだろうか。この不安定さはコンテスト中に使うのが厳しい……。

## 行列

https://doc.sagemath.org/html/en/reference/matrices/index.html

`matrix([[1, 2], [3, 4]])` として行列が作れる。 `matrix(GF(10^9+7), [[1, 2], [3, 4]])` なら `mod 1000000007` で計算。

https://atcoder.jp/contests/abc200/tasks/abc200_f

行列累乗で（も）解ける問題。

```abc200_f.sage
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。残念。

https://atcoder.jp/contests/abc200/submissions/45704644

## 文字列

https://doc.sagemath.org/html/en/reference/combinat/sage/combinat/words/finite_word.html

`Word("abc")` に文字列関連のアルゴリズムがある。

https://atcoder.jp/contests/abc320/tasks/abc320_b

最長の回文を探す問題。想定解法は単純な $O\left(\left|S\right|^3\right)$ のアルゴリズムだが、回文については色々とアルゴリズムがあり、SageMathにも実装されている。

```abc320_b.sage
print(max(Word(input()).lps_lengths()))
```

これは余裕で最短コードだろと思ったら、上にNibblesという言語がいた。コードゴルフ用の言語らしい。なるほど……。

## グラフ

https://doc.sagemath.org/html/en/reference/graphs/index.html

解説で名前が出てくるようなアルゴリズムはだいたいありそう。

https://atcoder.jp/contests/typical-algorithm/tasks/typical_algorithm_d

最短経路。

```typical-algorithm_d.sage
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"]))
```

## 多項式

https://doc.sagemath.org/html/ja/tutorial/tour_polynomial.html

https://atcoder.jp/contests/abc245/tasks/abc245_d

多項式の除算をそのまま書けば良い。`ZZ` は整数（$\mathbb{Z}$）上の多項式だということ。

```abc245_d.sage
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)
```

ラグランジュ補間もあるので、この問題がさくっと解けそう。

https://atcoder.jp/contests/abc208/tasks/abc208_f

https://atcoder.jp/contests/abc208/editorial/2195

……と思ったけれど、TLE。

```abc208_f.sage
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))
```

https://atcoder.jp/contests/abc208/submissions/45709779

## 数独

https://doc.sagemath.org/html/en/reference/games/sage/games/sudoku.html

ドキュメントを眺めていたら、数独があった。もし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|
+-----+-----+-----+
```