多倍長演算の活用②
Python の多倍長演算を活用する方法の第 $2$ 弾です。 前回の記事 では、多倍長整数の各bitをフラグと見て処理する方法を説明しました。この記事では、いよいよ整数を要素に持つ配列の演算を扱います。 ここからが本番です。 最後まで読んで頂けると嬉しいです。
FFTの代用としての畳み込み乗算への応用や、形式的べき級数との関係についても紹介します。
整数と配列の対応
多倍長整数を $k$ bitごとに区切って、それぞれが $k$ bit の非負整数を表していると見ることを考えます。なお文字 $k$ は、コード中も含めこの記事を通して常にこの意味で使います。
具体的には、 $0$ 以上 $2^k$ 未満の要素からなる配列 $A = [a_0,\ \ldots,\ a_{N-1}]$ について、対応する整数を
$${\rm number}(A) = \sum_{i=0}^{N-1}\ a_i \cdot 2^{ki} $$
と定めます。逆に整数 $n$ について、配列 ${\rm list}(n)$ を $${\rm list}(n)[i] = {\rm int}(n\ /\ (2^{ki}))\ {\rm mod}\ (2^{k}) $$ で定義します。 $A[i]$ は $A$ の $i$ 番目の要素を表します。
すると ${\rm number}$ と ${\rm list}$ は互いに逆写像になり、非負整数全体からなる集合と $0$ 以上 $2^k$ 未満の要素からなる配列全体からなる集合の間に 全単射 を与えます。ただし、配列 $A$ の末尾についている $0$ は無視し、これを外して一致するものは同一視することにします。
例として $A = [0, 1, 2, 3]$ を考えます。 $2$ 進法で書くと $A = [00_{(2)}, 01_{(2)}, 10_{(2)}, 11_{(2)}]$ ですね。例えば $k = 2$ としてこれを整数に変換すると ${\rm number}(A) = 228 = 11100100_{(2)}$ になります。定義により $A$ の小さい添え字のものが小さい桁に来るので、見た目の左右が逆になる点に注意してください（もちろん場合によっては左右を逆で持つ方がよい場合もあります）。
コードで書くとこんな感じです。
# 配列 -> 整数 (説明用の愚直なので遅い)
def nu(k, L):
s = 0
for i, a in enumerate(L):
s += a << i * k
return s
# 整数 -> 配列 (説明用の愚直なので遅い)
def li(k, n):
L = []
while n:
L.append(n & ((1 << k) - 1))
n >>= k
return L
# 本文で紹介したサンプル
K = 2
A = [0, 1, 2, 3]
n = nu(K, A)
print(n, bin(n)) # 228 0b11100100
print(li(K, n)) # [0, 1, 2, 3]
# オーバーフローして正しく復元できないケース
K = 2
A = [3, 14, 159, 265, 3589]
n = nu(K, A)
print(n, bin(n)) # 938347 0b11100101000101101011
print(li(K, n)) # [3, 2, 2, 1, 1, 0, 1, 1, 2, 3]
# Kを大きくしてオーバーフローを回避
K = 12
A = [3, 14, 159, 265, 3589]
n = nu(K, A)
print(n, bin(n)) # 1010231904743514115 0b111000000101000100001001000010011111000000001110000000000011
print(li(K, n)) # [3, 14, 159, 265, 3589]
演算
ここから演算の対応を考えます。やりたいことのイメージは、配列の世界だと演算をするのにたくさん（例えば O(N) ）の計算が必要な処理であっても、整数の世界で考えることで対応する演算を $O(1)$ 回の多倍長演算でできる場合があります。これを使って、配列を一度整数に置き換えることでコードの簡潔化や処理の高速化をしたいです。
各点和
配列 $A$ と配列 $B$ の和を
$$
(A + B)[i] = A[i] + B[i]
$$
で定めます。要素で書くと、 $A = [a_0,\ a_1,\ \ldots,\ a_{N-1}]$ と $B = [b_0,\ b_1,\ \ldots,\ b_{N-1}]$ について、 $A+B = [a_0+b_0,\ a_1+b_1,\ \ldots,\ a_{N-1}+b_{N-1}]$ です。 $A$ と $B$ の要素数が一致していない場合は、末尾に $0$ が付いていると思って定義します。普通にやると $O(N)$ の計算量が必要ですね。ここで、 $A+B$ のすべての要素が $2^k$ 未満である場合 、
$$
{\rm number}(A+B) = {\rm number}(A) + {\rm number}(B)
$$
が成立することが分かります。つまり、（左辺の）配列の世界では $N$ 回の演算をが必要だった計算が、（右辺の）整数の世界では $kN$ bit の整数演算を $1$ 回行うだけで実現できます 。
イメージしやすいように10進法で説明すると、 $A = [1,2,3]$ と $B = [3,4,5]$ の和を計算すると、配列で見ると3回の演算で $[4, 6, 8]$ が得られるのに対し、 $321 + 543 = 864$ とやると1回の演算で計算できる感じです。ただし $A = [1,2,3]$ と $B = [3,4,8]$ の場合、 $321 + 843 = 1164$ のように 「繰り上がり」（オーバーフロー） が発生してしまうため正しく復元できません。なおこの場合は2桁ずつ区切って $30201 + 80403 = 110604$ と見ることで正しく復元できます。$2$ 進法でも同様です。 $k$ は途中で繰り上がりが発生しない範囲でなるべく小さくするのが効率が良いです。
オーバーフローについては、以下でも同様に注意が必要です。
スカラー倍
配列 $A$ の各要素を定数倍（ここでは $t$ 倍）することを考えます。すなわち、
$$
B[i] = t * A[i]
$$
です。これは整数の世界でやると下記になります。$B$ のすべての要素が $2^k$ 未満である範囲において
$$
{\rm number}(B) = t * {\rm number}(A)
$$
が成立します。
10進法で書くと、 $A = [1, 2, 3]$ について、 $3 * A = [3, 6, 9]$ になるのは $3 * 321 = 963$ に対応します。ここでも「繰り上がり」には注意してください。
畳み込み
配列 $A$ と配列 $B$ の畳み込みを
$$
{\rm conv}(A, B)[i] = \sum_{j=0}^{i} A[j] * B[i-j]
$$
で定義します。
配列の畳み込みは整数では掛け算に相当します。 ${\rm conv}(A, B)$ のすべての要素が $2^k$ 未満である範囲で
$$
{\rm number}({\rm conv}(A, B)) = {\rm number}(A) * {\rm number}(B)
$$
畳み込み演算は愚直にやると $O(N^2)$ （FFTを使うと $O(N\log N)$）ですが、これが1回の乗算でできることになります。
これも10進法で説明すると、 $A = [1,1,1]$ と $B = [1,2,3]$ の畳み込みは $[1, 3, 6, 5, 3]$ ですが、整数では $111 * 321 = 35631$ になります。
この畳み込みの考え方はとても重要です。 下で紹介する問題でも何度か出てきます。
合計
配列 $A$ を表した整数 $x$ が与えられたときに、 $A$ の合計を計算する方法です。 $l$ を $A$ の要素数の bit 数とすると、下記のように $O(l) = O(\log N)$ 回の bit 演算で計算できます。ただし、ここでも「繰り上がり」には注意してください。前回の記事で紹介した popcount を求める手法に似ていますね。
def getsum(k, x):
for i in range(l):
x += x >> (k << i)
return x & ((1 << k) - 1)
数え上げ問題
次の問題を考えます。
$N$ 個の荷物があり、 $i$ 番目の荷物の重さは $w_i$ です。これらから重さの合計がちょうど $W$ となるようにいくつかの荷物を選ぶ方法は何通りあるか。
前回の記事 では、可能な重さをすべて列挙する方法を扱いましたが、今回は可能かどうかだけでなく「何通りあるか」の情報を持たせる必要があります。
配列を用いる方法
$j$ 番目まで見て重さがちょうど $i$ になる選び方を $X[j][i]$ とすると $X[j+1][i] = X[j][i] + X[j][i-w_j]$ みたいに $DP$ できます。第一引数をなくして使いまわすことにより $1$ 次元にすると $X[i] = X[i] + X[i-w_j]$ みたいに推移させればよいです。計算量は $O(NW)$ です。
A = [1, 1, 1, 2, 2, 3, 4, 6] # 荷物の重さ
W = 10 # 重さの最大値
X = [0] * (W+1)
X[0] = 1 # 最初は重さ0だけ可能
for a in A:
for i in range(a, W+1)[::-1]: # 同じ荷物を2回使わないように、大きい方から更新する
X[i] += X[i-a]
print(X)
# [1, 2, 3, 5, 6, 7, 9, 10, 10, 11, 11]
多倍長演算を用いる方法
これを整数で表現することを考えましょう。推移は s *= (1 << a * K) + 1 あるいは s += s << a * K などのように簡潔に書くことができます。 $O(N)$ 回の bit 演算で処理できます。
A = [1, 1, 1, 2, 2, 3, 4, 6] # 荷物の重さ
s = 1
for a in A:
s += s << a * K
print(li(K, s))
# [1, 3, 5, 8, 11, 13, 16, 19, 20, 21, 22, 21, 20, 19, 16, 13, 11, 8, 5, 3, 1]
なおこれだと重さが $W$ を超えた分も情報を保持してしまって無駄になってしまうので、 前回 と同様、マスクをかけるか逆順にすることによって不要な部分を保持しないようにすることもできます。
各点 MOD
上のようなナップサック問題などでは、結果が大きくなるため ${\rm mod}\ P$ を取る必要があることが多いです。処理を $1$ つするたびに毎回整数を配列に変換して $\rm mod$ を取りまた整数に戻す、のような処理をするととても非効率です。ここでは、多倍長整数を $k$ bit ごとに ${\rm mod}\ P$ を取る方法を説明します。
配列 $A$ の各点を ${\rm mod}\ P$ することを考えます。簡単のため、ここでは $P$ は $30$ bit程度の素数とします（実は素数じゃなくても奇数なら大丈夫です）。これの $3$ 倍に少し余裕をもって $k = 96$ としておきます。
毎回の乗算などで要素の桁数は $30$ bitぐらい増えてしまいますが、これが繰り上がりにならないためには、毎回 $64$ 桁ぐらいに抑えられればよいです。 $96$ bitを上位 $32$ bit と下位 $64$ bitに分けておきます。
事前に $t = 2^{64}\ {\rm mod}\ P$ を求めておくことで、上位 $32$ bitを下位 $64$ bitに持っていくことができます。具体的には
$$
a * 2^{64} + b \equiv t\cdot a + b
$$
のようにします。 $t$ は 30bit、 $a$ は $32$ bit、 $b$ は $64$ bitなので、右辺は最悪でも $65$ bitに収まります。すると、次に $30$ bit程度の掛け算をしても $96$ bitに収まるのでオーバーフローの心配はありません。上位 bit と下位 bit に分けるのは、適当なマスクをかければできます。具体的には、 $2$ 進法で $111\cdots111000\cdots000$ （$1$ が $32$ 個と $0$ が $64$ 個並ぶ）をたくさん並べてできる整数と $\rm bitwise\ and$ すると、各 $k$ bit ごとの上位 $32$ 桁だけが残ります。
この考え方も、多倍長演算での有用性を高めることにつながっています。これにより 「答えが非常に大きくなるため $\rm mod$ を取らないといけない問題」も自然に多倍長演算で実装できる ことになります。
後ほど扱う問題でも出てきます。
例題
高速フーリエ変換の代用1
「畳み込みをするだけ」ですね。タイトルを見てもいかにも高速フーリエ変換（ FFT ）で $O(N\log N)$ を想定している感じですが、多倍長乗算 $1$ 回（と $O(N)$ 回の通常の演算）で実装できます。
簡単にコードの解説をします。このラムダ式 $3$ つが肝ですね。
nu = lambda L: int("".join([bin(K+a)[-k:] for a in L[::-1]]), 2)
st = lambda n: bin(n)[2:] + "0"
li = lambda s: [int(a, 2) if len(a) else 0 for a in [s[-(i+1)*k-1:-i*k-1] for i in range(len(A)+len(B)-1)]]
記事の冒頭で説明した「愚直」の nu と li を多少高速化しています。コード中の st は「整数 $\rightarrow$ 文字列」 、 li は「文字列 $\rightarrow$ 配列」です。途中に文字列をかませて高速化しています。 （Python では、 $2$ 進、 $4$ 進、 $16$ 進においては、文字列と整数の変換は比較的高速にしてくれます。）
nu は「配列 $\rightarrow$ 整数」の変換です。ここでも途中、文字列を経由しています。 なおこのあたりの変換は経験的にこのACコードの方法（あるいはこれを $16$ 進にしたもの）が比較的速い気がしていますが、実際のところどれが一番高速かは確信が持てていません。 Python と PyPy でも違うと思います。もっと速い方法やアイデアがあれば教えてください。
実質的な実装は下記の部分です。
for a in li(st(nu(A) * nu(B)))[1:]:
print(a)
やっていることのイメージはこんな感じです：
- まず配列 $A$ と $B$ を
nu(A)とnu(B)で整数化します。 - 次に、それらの（整数としての）積を計算します。
- それを
stおよびliで配列に戻します。 - これらにより、配列 $A$ と $B$ の畳み込みを計算したことになります。
高速フーリエ変換の代用2
小さい $N$ の畳み込みをたくさんやるパターンです。こういう設定なら FFT より速い です！
高速フーリエ変換の代用3
非想定ですが、これもできます。
高速フーリエ変換の代用4
こちらはAGCから。積の畳み込みですね。
その他AGCの過去問
解説解法にのっとって、 ($1$ の個数 - $2$ の個数) ごとに集計することを考えます。これを多倍長整数で表します（負のものもあるので適当にシフトしてずらします）。遷移は直近の $3$ 項だけから決まる（ $4$ 項間漸化式のようになっている）ので、直近の $3$ 項に対応する多倍長整数を持てば処理できます。
遷移の本質はこの $2$ 行です。 $a$ 、 $b$ 、 $c$ はそれぞれ $1$ つ前、 $2$ つ前、 $3$ つ前の状態を表します。
for i in range(1, 3 * N + 1):
a, b, c = modP(((a << K) + (b >> K) + c) * pow(i, P-2, P)), a, b
形式的べき級数を使う問題
ABC159-F Knapsack for All Segments
ACコード
考察は maspyさんの記事 が参考になります。記事では形式的べき級数を用いて説明されていますが、書かれているとおり、
$$
F_i = (F_{i-1} + 1)(1+x^{A_i})
$$
みたいな推移をする必要があります。これは多倍長整数の書き方では
s = (s + 1) * ((1 << A_i * k) + 1)
あるいは
s += 1
s += s << A_i * k
などのように記述できます。
ACコード では、シフトの左右を逆にして同様の処理をしています。
なお答えは非常に大きくなることがあるので、適宜 $\rm mod$ を取る必要があります。このコードでは、 $15$ 行目の s -= ((s & m) >> 30) * pa で各点 ${\rm mod}$ を取っています。
(1 + ai X) の総積
$(1+a_0x)(1+a_1x)\cdots(1+a_{N-1}x)$ のような式を展開する問題がよくありますね。 （下に書きますが、ここ3日で類題4問ほど見ましたよ！）
これも多倍長演算ととても相性が良いです。いくつか例を挙げて説明します。Easy と Hard の大きく2種類の実装を紹介します。
(1 + ai X) の総積 Easy 1
NOMURA2020-D: Urban Planning
ACコード
$(1+a_0x)(1+a_1x)\cdots(1+a_{N-1}x)$ を展開することを考えます。愚直にやると $O(N^2)$ ですが、これを高速化することを考えます。
ここまで読んだみなさんなら簡単ですね。順にかけていくと O(N) 回の bit 演算でできます。 modP を事前に定義しておけば、遷移は実質2行です。
s = 1
for a in A:
s *= (a << k) + 1
s = modP(s)
(1 + ai X) の総積 Easy 2
ABC169-F Knapsack for All Subsets
ACコード
タイトルとは少し違いますが、 $(2 + X^{a_1})(2 + X^{a_2})\cdots (2 + X^{a_N})$ の $X^S$ の係数を求めれば良いです。$X^S$ より大きい次数のところはいらないので、「左右逆」にして右側にシフトするようにすればいらないところがカットされます。
遷移はこんな感じ。
s = 1 << k * S # S番目からはじめて右にシフトすることで、不要な部分をカットできる
for a in A:
s = 2 * s + (s >> k * a)
s = modP(s)
(1 + ai X) の総積 Easy 3
Yukico250-D いろいろな色 (Easy)
ACコード
ところで、この Easy 1, Easy 2, Easy 3 で取り上げている3問、なんと3日連続で出たんですね。
これは重要項目として是非マスターしましょう。
(1 + ai X) の総積 Hard
Easy の方法では $N$ が大きく（5～10万程度）なるときついこともあります。この場合はかける順番を工夫するとさらに高速化できます。
アイデアとしては、1つずつ順番にかけていくと後半ではとても桁数の多いものとそんなに桁数が多くないものをかけるので、効率が悪いです。なのでなるべく桁数の大きい計算はまとめてやるために、同じ大きさのものをまとめてやると良いです。
これはセグ木の初期化の要領で実装できます。具体的には、まず添え字が $N$ から $2N-1$ の範囲にもとの配列を入れておいて、 $i = N-1,\ \cdots,\ 1$ について $A[i] = A[2i] * A[2i+1]$ をすれば、すべての総積が $A[1]$ に格納できます。
for i in range(N)[::-1]:
A[i] = modP(A[2*i] * A[2*i+1])
なお、逐次で掛け算する場合に比べて桁数が多めに必要なので modP の定義を少し変えています。少なくとも2回は mask と掛け算の処理をしないといけないです（下記では3回やっています）。 m1 などは最初から $5000 * K$ 桁持たせていますが、少しずつ増やした方が速い気もしています。
P = 10 ** 9 + 7
k = 12
K = k * 8
pa1 = (1 << k * 4 + 16) - ((1 << k * 4 + 16) % P)
pa2 = (1 << k * 2 + 24) - ((1 << k * 2 + 24) % P)
pa3 = (1 << k + 28) - ((1 << k + 28) % P)
m1 = int(("1" * (k * 4 - 16) + "0" * (k * 4 + 16)) * 5050, 2)
m2 = int(("1" * (k * 6 - 24) + "0" * (k * 2 + 24)) * 5050, 2)
m3 = int(("1" * (k * 7 - 28) + "0" * (k + 28)) * 5050, 2)
def modP(x):
x -= ((x & m1) >> k * 4 + 16) * pa1
x -= ((x & m2) >> k * 2 + 24) * pa2
x -= ((x & m3) >> k + 28) * pa3
return x
ABC169-B Multiplication 2
ACコード
mod を取らずにただただ総積を計算する問題です。（非想定ですが、これで場合分けがいらなくなります。）
ところでこの問題も Easy 2 と同じ日ですね！
Yukico250-D いろいろな色 (Hard)
筆者はまだこの方法では通せてないです。。誰か通してください。
多項式・形式的べき級数と多倍長演算の関係
上の例でもあったように、実は 多項式・形式的べき級数と多倍長演算はとても相性が良いです 。これは 多項式・形式的べき級数できれいに書ける演算の多くは、多倍長整数の演算でもきれいに書ける ということを意味します。ただし、無限に続く形式的べき級数は整数に対応させることができないので、ここでは有限の多項式のみを考えます。多項式は（係数を配列と見ることで）配列の場合と同様に整数に対応させることができます。
多項式 $f$ および $g$ に対応する整数をそれぞれ $s$ および $t$ とすると、次のような関係にあります。すなわち、下記の形式的べき級数の演算はいずれも $O(1)$ 回の多倍長演算で実装できます。
| 形式的べき級数での演算 | 多倍長演算での演算 | |
|---|---|---|
| 加算 | $f + g$ | $s + t$ |
| 定数倍 | $mf$ | $ms$ |
| 畳み込み乗算 | $f\cdot g$ | $s * t$ |
| 文字式をかける | $x^a \cdot f$ | $s << a * k$ |
| 文字式で割る | $\displaystyle \frac{f}{x^a}$ | $s >> a * k$ |
「左右逆」の実装をすると、 $1+x$ での割り算なども考えることができます。ただし無限に続く分は処理できないので、予め設定した桁数で打ち切られます。
| 形式的べき級数での演算 | 多倍長演算での処理 | |
|---|---|---|
| 割り算 | $\displaystyle \frac{f}{1-x} \ = (1 + x + x^2 + \cdots)f$ | ($s<<k)\ //\ ((1 << k) - 1)$ |
| 割り算2 | $\displaystyle \frac{f}{1-x^a} \ = (1 + x^a + x^{2a} + \cdots)f$ | ($s<<k)\ //\ ((1 << k\ *\ a) - 1)$ |
もちろん、多項式・形式的べき級数でかけても多倍長演算で扱いにくい問題や、その逆のケースもあると思います。ただ両者はある程度共通していることは確かだと思います。これは、いずれも 数学的にきれいな処理 を扱っているからだと思っています （筆者の主観） 。
多倍長演算でできない処理
各点積・各点max
本記事では、各点和・各点 mod ・畳み込み などを扱いました。 一方、多倍長演算では扱いにくい演算として、各点積、各点 max 、各点 min などがあります。 （筆者が思いついていないだけの可能性もあります。いい方法があれば教えてください。）
→ できるみたいです！下に書きます。
ナップサック問題
実は、この記事の執筆中に solzard さんから難しいお題を頂きました。
重さと価値が与えられる一般的なナップサックでも、上手くやればビット演算で解けると思いますか?
— solzard (@solzard_) March 29, 2020
多倍長の記事はほとんどないので、続編も楽しみです
実は難しい気がしています（笑）
— きり (@kiri8128) March 30, 2020
理由はいくつかあるのですが、各点 max みたいなのを取る必要があり、それが bit 演算だけでは難しい気がしています（私が思い付いてないだけかもですが）。もし良い方法があれば教えてください。
結論を言うと、効率的に計算するのは難しい気がしています。ツイートでは各点 max を取得するのが難しいと書きましたが、それ以外にも難しいポイントはたくさんあります。そもそも重さと価値を同時にきれいに扱える方法がなかなか思いつきません。。
と思いきや、くりんぺっとさんからアイデアを頂きました！！
早いかどうかはさておき、min(a,b) = b + ((a-b) & (-1 if a-b<0 else 0)) と変形できて、条件式の部分はa-bの符号ビットを右シフトしてスカラー倍で表現できます。繰り下がりが起きないよう2^{k-1}を足しておくなどの工夫は必要ですが、各点min/maxは実現できると思います。
— くりんぺっと (@climpet) June 2, 2020
ってことで実装しました。
ポイントは各点 max と遷移の式です。
各点 max
def MAX(a, b):
return b + (((((a + ms - b) & ms) >> (K - 1)) * tt) & (ms + a - b))
ツイートで書かれているのをそのまま実装しています。
ms は各点 INF （一番上の桁に1がある）です。
遷移
s = 1 << K - 2 + (W * K)
for _ in range(N):
w, v = map(int, input().split())
s = MAX(s, (s >> w * K) + one * v)
$O(1)$ 回の bit 演算で遷移を表現できています。
ただまだスピード的にはあまり高速とは言えないですね。。多倍長の乗算を使っているところで時間がかかっているんでしょうか。このへんは今後の課題としたいと思います。
終わりに
$2$ 回にわたって、多倍長整数を使って処理を簡潔化・高速化する話をしました。 Python で競技プログラミングをするのであれば、紹介した手法が使える機会はそこそこあるかと思います。ただ当然ながらこれらの方法がすべての場合に適用できる訳ではありません。それはそうで、みんな大好き（？）ダイクストラ法でもすべての問題で役立つかというとそうでもないと思います。ダイクストラ法も多倍長演算による方法も、あくまで $1$ つの手法として「使える場合には使える」程度に捉えてもらえれば良いかと思います。
ただ個人的にはせっかく多倍長演算が使える Python を使うんだから、使える問題ではぜひ使っていきたいという気がします。なんでかって、数万桁の乗算をするなんてテンション上がりませんか？私は上がります。
2020年3月31日：投稿
2020年5月31日： (1 + ai X) の総積に関する記事を記載
2020年6月3日： ナップサック問題について追記。あと（ネタバレ防止のため伏せていたけど）例の問題名を記載した方が分かりやすい気がしたので記載