多倍長演算の活用
Python の多倍長演算を活用して、処理の簡潔化・高速化を行う方法について書きます。
本記事では整数の各bitをフラグとみてbit演算する処理について、 次の記事 では整数を要素とする配列の処理について書きます。
後半、一部ネタバレを含むのでご注意ください。
ナップサック問題
例えば、次のような問題を考えます。
$N$ 個の荷物があり、 $i$ 番目の荷物の重さは $w_i$ です。これらから重さの合計が $W$ 以下となるようにいくつかの荷物を選ぶとき、重さの最大値を求めよ。
この問題を解くには、順番に荷物を見て、その時点でどの重さにできるかのリストを更新できれば良いです。
配列を用いる方法
配列でこれを表す場合、大きさ $W+1$ の配列 $X = [x_0,\ \dots,\ x_{W}]$ を用意して、荷物を前から順に見て、その時点で荷物で重さ $i$ にできるのであれば $X_i = 1$ 、そうでなければ $X_i = 0$ とすることで $O(W * N)$ 回の処理でできます。配列は使いまわせば $1$ 次元で大丈夫です（ $X[j][i]$ を $j$ 番目までの荷物を見たときに $i$ にできるか、とみて $2$ 次元配列にすることもできます）。
A = [3, 5, 7, 11] # 荷物の重さ
W = 20 # 重さの最大値
X = [0] * (W+1)
X[0] = 1 # 最初は重さ0だけ可能
for a in A:
for i in range(a, W+1)[::-1]: # 同じ荷物を2回使わないように、大きい方から更新する
if X[i-a] == 1:
X[i] = 1
print(X)
# [1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 0]
bit演算を用いる方法
bit演算で行う場合は、重さ $i$ にできる場合には $s$ の下から $i$ 番目のbitを $1$ にして、そうでなければ $s$ の下から $i$ 番目のbitを $0$ にします。重さ $a$ の荷物の更新処理は、 $s$ に、 $s$ を $a$ だけ左シフトしたものと ${\rm bitwize}\ or$ を取れば良いです。この方法だと $O(N)$ 回のbit演算で解くことができます。wordsizeを $w$ として $a$ bitのbit演算の計算時間が $O(a/w)$ でできるとすると全体の計算時間は $O(W * N / w)$ です。
A = [3, 5, 7, 11] # 荷物の重さ
s = 1 # 最初は0番目のbitだけ立てる
for a in A:
s |= s << a
print(bin(s))
# 0b100101011011101110110101001
最後の結果は、配列の方法と左右反転していることに注意してください。ただし、この方法では、 $W$ を超えた部分も計算してしまっているところが無駄になっています。これをやめたければ、毎回 $m = (1<<W) - 1$ と ${\rm bitwize}\ and$ を取れば良いです。もう一つの方法としては、最初 $W$ 番目のbitを立てておき、更新を右シフトで行う方法もあります。これなら不要な部分は勝手に削除されるので高速です。
A = [3, 5, 7, 11] # 荷物の重さ
W = 20
s = 1 << W # 最初はW番目のbitだけ立てる
for a in A:
s |= s >> a
print(bin(s))
# 0b100101011011101110110
AGCの過去問
（ネタバレ防止のためリンク先は記載していません。リンク先に飛んで確認してください。）
こちらの 問題 をbit演算で解くことを考えます。
本問では、入力を適当に2進法で受け取った後、 ① 二項係数テーブルの計算、 ② 与えられた数値との bitwize and 、 ③ それらのすべての xor の取得、をすべてbit演算で行うことができます。先にいくつか準備します。
二項係数 mod 2
maspyさんのアイデアをお借りして、bit演算で $_NC_k\ {\rm mod}\ 2\ (k = 0,\ 1,\ \dots,\ N)$ のテーブルを $O(\log N)$ 回のbit演算で求めます。
これすごい。
— きり (@kiri8128) March 22, 2020
nCkが奇数となるすべてのkをlogn回のbit演算で計算できるのか。 https://t.co/qOOpHy2H9J
要は $N$ の下から $i$ 番目のbitが立っていれば、 s |= s << i すれば良いですね。 $N$ のLSB（least significant bit）は N & (-N) で求められるので、下記のようにできます。
N = 9
s = 1
while N:
i = N & -N
s |= s << i
N ^= i
print(bin(s))
# 0b1100000011
全ビットの xor
$k$ bitの整数 $N$ の、すべてのビットのxorを取ります。これは $O(\log k) = O(\log\log N)$ 回のbit演算で可能です。具体的には、popcountを計算する要領で、各桁の情報を集約させていきます。popcountの取得のアイデアは、いろんな記事（例えば こちら ）で書かれているので適当に調べてください。
桁数が大きくなる場合（数万 bit になる場合）のpopcountは下の「類題」でも出てきます。
xor を取るだけなら、マスクなどがいらない分、簡潔になります。
def popcount_parity(n):
for i in range(n.bit_length().bit_length()):
n ^= n >> (1 << i)
return n & 1
AC コード（参考）
（ちょっとコードが汚いですが）これを使った ACコード です。他の Python の AC に比べると、ずいぶん高速になっていることが分かるかと思います。
その他の類題
（ネタバレ防止のため、リンクのみの記載としています。リンクに飛んで確認してください。）
たくさんある気がしたけど思いつかない（他にもあったら教えてください）。
類題 1
まさに本記事で扱った方法で処理できます。最後の popcount を取るところも bit 演算で処理すると高速です（この制約なら $1$ 桁ずつ見ても一瞬ですが）。
類題 2
こちらも、最初の考察の後は似たような処理でできます。
類題 3
10000 bit 程度保存する必要がありますが、十分高速に処理できます。
$x$ の $i$ 番目のビットが立っているかを確認する際に、 x & (1<<i) > 0 または (x>>i) & 1 > 0 のようなことをやりますが、これを桁数分だけやると $2$ 乗になりそうなので （未確認） 、その場合は一度 s = bin(x) のように文字列にしてしまえばその後の判定は各 $O(1)$ でできます。
（追記）よく考えればこれも LSB と bit_length() 、あるいは popcount の二分探索などでやっても良かったですね。
類題 4
（追記）4と5は本記事で解説したシフト演算などをほとんど使いませんが、並行処理により高速化できる例として挙げておきます。
いかにも bit 演算を使いたくなる形をしています。 $N$ bit の整数で path にどの頂点が含まれるかを表すことができます。複数の path を合体させる際は単純に ${\rm bitwise}\ or$ を取ればよいのできれいですね。
bit DP みたいにするところがやや複雑なので、下記コード中にコメントで説明します。
for m in range(1<<M):
a = m & (-m) # LSB（これをmとxorすれば、bitを1つ消したものが得られる）
if a == m:
if a == 0:
# bit が 1 つも立っていないとき
Z[m] = 0
else:
# bit が 1 つだけ立っているとき
Z[m] = Y[D[a]]
else:
# bit が複数立っている場合は、1つ減らしたものと or を取る
Z[m] = Z[m^a] | Y[D[a]]
# bitが立っていない分だけ自由度があるので、自由度ごとに集計
# 包徐原理により、popcountの偶奇でプラスマイナスが決まる（後半のpopcntはparityだけで良かった）
CC[N - 1 - popcnt(Z[m])] += (1 if popcnt(m) & 1 == 0 else -1)
類題 5
これは想定が $2000$ の $3$ 乗なので Python だと多倍長を使わないときつい気がします。
マニアックな方法
元ネタはこどふぉの 問題 です。ややマニアックですが、工夫すれば bit 演算でいろんなことができるというのを知っていただきたいので説明します。
ポイントだけ説明するとこんな感じです。
$\unicode{x201C}S\unicode{x201D}, \unicode{x201C}E\unicode{x201D}, \unicode{x201C}T\unicode{x201D}$ からなる長さ $K$ の文字列が $N$ 個あります。 $\unicode{x201C}S\unicode{x201D}, \unicode{x201C}E\unicode{x201D}, \unicode{x201C}T\unicode{x201D}$ からなる長さ $K$ の文字列 $A$ 、 $B$ について、 $C = f(A, B)$ を
① $A[i] = B[i]$ なら $C[i] = A[i]$
② $A[i] != B[i]$ なら $C[i] != A[i]$ かつ $C[i] != B[i]$
で定めます（ $A[i]$ は $A$ の $i$ 番目の文字を表します）。例えば、 $A = \unicode{x201C}SSE\unicode{x201D}$ かつ $B = \unicode{x201C}SET\unicode{x201D}$ なら $C=\unicode{x201C}STS\unicode{x201D}$ です。
$N$ 個の文字列の異なる $3$ 個 $A$ 、 $B$ 、 $C$ の選び方で、 $f(A, B) = C$ を満たすものの個数を求めよ。
実は適当に文字列を整数に変換しておくと、 $f(A, B)$ が $O(1)$ 回の多倍長演算で計算できます。その方法を説明します。
3種類の文字があるので、$1$ bitで $1$ 文字を表すことはできません。 $2$ bitあれば情報を持たせることはできますが、その後の演算が自然に行えないので、ここでは $7$ bitで $1$ 文字を表すことにします。具体的には、 $\unicode{x201C}S\unicode{x201D}$ を $0000100$ 、$\unicode{x201C}E\unicode{x201D}$ を $0001000$ 、 $\unicode{x201C}T\unicode{x201D}$ を $0010000$ で表します。左右 $2$ bit ずつはダミーで必ずゼロになりますが、後の計算のために追加しています。
結論を言うと、
$$
f(A, B) = (m - (A\ {\rm or}\ B))\ {\rm xor}\ ((((A\ {\rm and}\ B) >> 2) * 31) \ {\rm and}\ m)
$$
で計算できます。ただし、 $m$ は $2$ 進表記で $0011100$ を $K$ 個並べた整数です。
1つ目のカッコ内は、 $x$ と $y$ が異なるとき、 「 $x$ でも $y$ でもないやつ」を表してくれるのでOKです（今思えば - は ^ でも良かった）。
ただし $x$ と $y$ が一致するときは $111\ {\rm xor}\ x$ になってしまうので、小細工して全体に $111$ を $xor$ したいです。具体的には、(x & y) >> 2 に 31 をかけると $x$ と $y$ が一致するときのみ $1$ がたくさん出てきます。左右にはみ出てしまう可能性があるので、最後にマスクすればOKです。このはみ出た分が邪魔しないようにダミーを付けたのでした。
N, K = map(int, input().split())
X = [int(input().replace("S", "0000100").replace("E", "0001000").replace("T", "0010000"), 2) for _ in range(N)]
D = {x: 1 for x in X}
m = int("0011100" * K, 2)
ans = 0
for i in range(N):
for j in range(i+1, N):
if (m - (X[i] | X[j])) ^ (((X[i] & X[j]) >> 2) * 31 & m) in D: ans += 1
print(ans // 3)
こんな感じで、工夫をすると bit 演算だけでいろいろできることが分かるかと思います。ただし複雑になりすぎると余計時間がかかることもあるので、並行処理のメリット（ざっくり wordsize 倍ぐらいの高速化？）と比較して適当に判断してください。
もし他にも面白いことができるよ、というご意見あれば教えてください。
次回は整数を要素とする配列の bit 演算による処理について説明します。
（2020年3月29日）投稿
（2020年3月30日）類題4と5を追加
（2020年3月31日）次の記事 を投稿