【ARC148 F】998244353 → 1000000007 別解（銀diff 3621）

筆者はこの問題を銀diffの中でも簡単な方だと考えています。自力で解きたい方はここで引き返してください！

問題概要

$0$ 以上 $1000000007$ 未満の整数 $a, b$ が与えられるので、$a \times b \mod 1000000007$ を計算するプログラムを書く。

ただし使える演算は以下の3つのみ。

符号なし64bit整数での加算（結果は $\mod 2^{64}$）
符号なし64bit整数での乗算（結果は $\mod 2^{64}$）
$998244353$ を法とする剰余演算

制約

命令数 $N$ は $1 \leq N \leq 100$

考察

大まかな方針

まず素直に考えると、$\mod 1000000007$ の演算が直接できないので詰まりそうに見える。引き算・割り算は逆元を使えばどうにかなるかも……と考えると、こんなアイデアが浮かぶ：

「$A \times B$ から $1000000007$ の倍数をうまく引き続ければ、余りが求まるのでは？」

$A \times B$ を $1000000007$ で割ったときの商を $D$、余りを $C$ とおくと、

$$A \times B = C + 1000000007 \times D$$

が成り立つ。もし $D$ が求まれば $C = A \times B - 1000000007 \times D$ として余りが計算できる。

$D$ を正確に求めるのは難しいので、発想を変えて 「$1000000007$ の倍数を少しずつ引いていく」 方針を取る。

具体例で確認

たとえば $12345 \mod 97$ を求める場合：

$$12345 - 97 \times 123 = 414$$
$$414 - 97 \times 4 = 26$$

引ける限り引き続けることで余りが求まる。引ける条件は「引かれる値 $\geq$ 引く値」であること。

ここで $97 < 100$ なので、$97k$ の代わりに $100k$（上位桁の値）を使って判定できる。たとえば $12345$ の上位3桁が $123$ なので $12345 \geq 12300 > 97 \times 123$ がわかり、$97 \times 123$ を引いて安全。この「上位桁を見て引ける量を決める」発想が基本方針となる。

2進数で上位ビットを取り出す

本問題では10進数の「桁」を直接扱えないので、2進数で考える。

$1000000007 < 2^{30} = 1073741824$ なので、

$$2^{30} \times k > 1000000007 \times k$$

が成り立つ。したがって、値の31ビット目以上の部分 $k = \lfloor X / 2^{30} \rfloor$ を求めれば、$1000000007 \times k$ を安全に引ける。

では、この $k$（上位ビット）をどうやって取り出すか。加算・乗算の結果が $\mod 2^{64}$ になる（オーバーフロー）ことを逆に利用する。

値 $X$ を $X = a \cdot 2^N + b$（$b < 2^N$）と表したとき、$X$ に $2^{64-N}$ を掛けると、

$$X \cdot 2^{64-N} = a \cdot 2^{64} + b \cdot 2^{64-N}$$

$a \cdot 2^{64}$ は $\mod 2^{64}$ でオーバーフローして消えるので、$b \cdot 2^{64-N}$ だけが残る。これに $\mod 998244353$ を取り、さらに $2^{64-N}$ の逆元 $\pmod{998244353}$ を掛けて再び $\mod 998244353$ を取れば $b = X \bmod 2^N$ が求まる。つまり、$X$ の下位 $N$ ビットが得られる。

イメージ図

上位ビット（$\lfloor X / 2^N \rfloor$）は $X - b$ を $2^N$ で割ることで求まる。

ただしこの操作には制約があり、$b < 998244353$ である必要があるため、$N \leq 29$ に限定される。

64ビット全体を扱う：2段階に分ける

上位 $K$ ビット（$K > 30$）を取り出したい場合、$N \leq 29$ の制約があるためそのままでは対応できない。そこで64ビットを複数のブロックに分けて処理する。

具体的には、64ビットを「上位 $(64-K)$ ビット」「中位 $(K-20)$ ビット」「下位 $20$ ビット」の3ブロックに分ける。

下位20ビットを求めて引く（ゼロにする）
$(64-K)$ ビット左にシフトして、中位 $(K-20)$ ビット部分を求めて引く（ゼロにする）
残った上位 $(64-K)$ ビットを $\mod 998244353$ で取り出す

この手順で、任意のビット幅の上位ビットを取り出せる。

詳細イメージ

ここまでの内容を踏まえて、具体的な目標を整理する。

$a \times b$ を計算する（$a, b < 10^9$ なので積は $2^{60}$ 未満）
$1000000007$ の倍数を繰り返し引いて、$2^{30}$ 未満にする
さらに引いて $1000000007$ 未満にする（＝余りが求まる）

しかしこの方針には2つの壁がある。

壁1：命令数が足りない

$2^{30}$ と $1000000007$ の差は $73741817$ と意外と大きい。毎回 $k = \lfloor X / 2^{30} \rfloor$ を求めて $1000000007 \times k$ を引く操作だと、3ビットずつしか減らせず命令数が100行を超えてしまう。

なぜ3ビットずつしか減らせないのかを説明する。

$1000000007$ を2進数で表すと 111011100110101100101000000111 であり、上位3ビットが $111$ で始まっている。これは $1000000007 > 2^{29} \times 7$ であることを意味する（$7 = 111_{(2)}$）。

ここで「上位ビットを $k$ として $1000000007 \times k$ を引く」操作を考える。$k = \lfloor X / 2^{30} \rfloor$ なので、引く量は：

$$1000000007 \times k > 2^{29} \times 7 \times \frac{X}{2^{30}} = \frac{7}{8} X$$

つまり $X$ の $\frac{7}{8}$ より多く引くので、残る値は必ず $\frac{X}{8} = \frac{X}{2^3}$ 未満になる。これがちょうど3ビット分小さくなることに対応している。

$a \times b$ は最大 $2^{60}$ 未満で、目標は $2^{30}$ 未満なので、消すべきビット数は30ビット分。3ビットずつ消していくと最低でも $\lceil 30 / 3 \rceil = 10$ 回の「引く操作」が必要になる。各操作に複数の命令が必要なことを考えると、簡単に100行を超えてしまう。

解決策：1000000007の倍数を使う

1回の操作でより多くのビットを減らすには、上位ビットが多く $1$ で並ぶような $1000000007$ の倍数を使えばよい。たとえば $1000000007 \times 2199 = 2199000015393$ は2進数で 11111111111111110100111010110001000100001 と表され、上位16ビットが $1$ で連続しており、一気に16ビット削減できる。同様に他の倍数も調べると、それぞれ使えるビット幅の範囲で効率よく削減できるものが見つかる。

$a \times b < 2^{60}$ なので、段階的に以下のように減らしていける：

使う値	削減ビット数	値のビット幅
初期値	—	60ビット幅
$1000000007 \times 2199$	16ビット	44ビット幅
$1000000007 \times 137$	7ビット	37ビット幅
$1000000007 \times 17$	3ビット	34ビット幅
$1000000007 \times 2$	3ビット	31ビット幅
$1000000007 \times 1$	1ビット	30ビット幅以下

これで命令数をかなり節約できる。

壁2：1000000007 以上 2^30 未満で詰まる

上記の操作を繰り返すと、値が $[1000000007,\ 2^{30})$ の範囲に入ってしまうことがある。この範囲では $\lfloor X / 2^{30} \rfloor = 0$ になるため引けるものがなく、操作が止まってしまう。しかしこの範囲の値は $1000000007$ 以上なので、まだ $\mod 1000000007$ の余りとして正しくない。

解決策：1000000007 を1回引いてみる

$1000000007$ を引いてみると、

元の値が $\geq 1000000007$ → 引いた結果は正、上位34ビットは 0
元の値が $< 1000000007$ → 引いた結果はアンダーフロー、上位34ビットは 全部1

つまり引いた結果の上位34ビットが0か1かを見れば、元の値が $1000000007$ 以上かどうかがわかる。

元の値が $\geq 1000000007$ だった → そのまま（引いたまま）
元の値が $< 1000000007$ だった → $1000000007$ を足して元に戻す

この「足すか足さないか」は上位34ビットの値（0か1）と対応しているので、「上位34ビットの値 × $1000000007$ を加算する」という操作で表せる。上位34ビットが0なら何も足さず、1なら $1000000007$ を足し戻すことになる。

上位34ビットの取り出し方：1ビット判定の工夫

素直に実装するなら、31ビット目を取り出して判定すればよい。具体的には10ビットと20ビットの2段階に分けて下位30ビットをゼロにし、$2^{33}$ 倍して $\mod 998244353$ で戻すことで31ビット目だけを取り出せる。しかし筆者がこの方法で実装したところ命令数が100を超えてしまったため、別の方法を考えることにした。

そこで発想を切り替える。先ほどの方法では下位30ビットをすべてゼロにできたため、31ビット目だけをきれいに取り出せた。しかし下位29ビットのゼロ化であれば30ビットのときより作業工程が1つ減る引き換えに、30ビット目が残ってしまい判定に干渉する。そこで「30ビット目も含めて場合分けしてしまえばよい」と発想を切り替える。下位29ビットをゼロにした後、値として取り得るパターンは4つしかない：

パターン	上位34ビット	30ビット目
①	0	0
②	0	1
③	1（アンダーフロー値）	0
④	1（アンダーフロー値）	1

これらに $\mod 998244353$ を取った値は：

パターン	$\bmod 998244353$
①	$0$
②	$536870912$
③	$856554439$
④	$395180998$

この4パターンの $\mod 998244353$ の値を見ると、①と②は上位34ビットが0、③と④は上位34ビットが1である。もしこれらに何らかの定数を掛けて $\mod 998244353$ を取った結果の最下位ビットが、①②は0、③④は1になるような定数が存在すれば、最下位ビットを見るだけで上位34ビットがわかる。そこで定数 $X$ を1から全探索したところ、$X = 4$ で成立することがわかった：

パターン	$\times 4 \bmod 998244353$	最下位ビット	上位34ビット
①	$0$	0	0 ✓
②	$150994942$	0	0 ✓
③	$431484697$	1	1 ✓
④	$582479639$	1	1 ✓

最下位ビット（0または1）が上位34ビットと一致しているので、この値を使って「最下位ビット × $1000000007$ を加算する」操作で、元の値が $1000000007$ 未満だったときだけ足し戻すことができる。

実装

※適宜符号なし64bit整数に置き換えています

以下が実際のプログラムである。

まず命令数を宣言し、$C = A \times B$ を計算する。

99
mul D A B

60→44ビット：$1000000007 \times 2199$ を使って16ビット削減する。最初の6行でまず $\mod 2^{20}$ を求めて $a \times b$ から引く。次の6行で $\mod 2^{43}$ を求めてさらに引く（下位43ビットをゼロに）。最後の6行で上位21ビットを求め $1000000007 \times 2199$ を掛けて $a \times b$ から引くことで44ビットにする。

mul E D 17592186044416
rem E E
mul E E 56644
rem E E
mul E E 18446744073709551615
add E D E

mul F E 8388608
rem F F
mul F F 113288
rem F F
mul F F 18446744073708503040
add F E F

rem F F
mul F F 453152
rem F F
mul F F 18446741874709536223
add F D F

44→37ビット：$1000000007 \times 137$ を使って7ビット削減する。

mul G F 17592186044416
rem G G
mul G G 56644
rem G G
mul G G 18446744073709551615
add G F G
mul H G 134217728
rem H H
mul H H 499129257
rem H H
mul H H 18446744073708503040
add H G H
rem H H
mul H H 7250432
rem H H
mul H H 18446743936709550657
add H F H

37→34ビット：$1000000007 \times 17$ を使って3ビット削減する。

mul I H 17592186044416
rem I I
mul I I 56644
rem I I
mul I I 18446744073709551615
add I H I
mul J I 1073741824
rem J J
mul J J 935854966
rem J J
mul J J 18446744073708503040
add J I J
mul J J 58003456
rem J J
mul J J 18446744056709551497
add J H J

34→31ビット：$1000000007 \times 2$ を使って3ビット削減する。

mul K J 17592186044416
rem K K
mul K K 56644
rem K K
mul K K 18446744073709551615
add K J K
mul L K 8589934592
rem L L
mul L L 366542959
rem L L
mul L L 18446744073708503040
add L K L
mul L L 464027648
rem L L
mul L L 18446744071709551602
add L J L

31→30ビット：$1000000007 \times 1$ を使って1ビット削減する。

mul M L 17592186044416
rem M M
mul M M 56644
rem M M
mul M M 18446744073709551615
add M L M
mul N M 17179869184
rem N N
mul N N 682393656
rem N N
mul N N 18446744073708503040
add N M N
mul N N 928055296
rem N N
mul N N 18446744072709551609
add N L N

30ビット→余り：$1000000007$ を1回引き、4パターン判定で必要なら足し戻す。

add N N 18446744072709551609
mul O N 34359738368
rem O O
mul O O 29001728
rem O O
mul O O 18446744073709551615
add O N O
rem O O
mul O O 4
rem O O
mul O O 9223372036854775808
rem O O
mul O O 890394177
rem O O
mul O O 1000000007
add C N O

まとめ

以上の方針をすべて組み合わせることで、命令数 99行 で $a \times b \mod 1000000007$ を計算するプログラムが完成し、ACを達成できた。

ポイントをまとめると：

基本方針：$1000000007$ の倍数を繰り返し引くことで余りを求める
上位ビットの判定：$1000000007 < 2^{30}$ を利用し、2進数の上位ビットで引く量を決める
上位ビットの取り出し：$\mod 2^{64}$ のオーバーフローと $\mod 998244353$ の逆元を組み合わせる
命令数の節約：$1000000007 \times 2199$、$1000000007 \times 137$ などを使って1回で多くのビットを削る
境界ケースの処理：4パターンに分類し、定数 $X = 4$ を掛けて $\mod 998244353$ を取った結果の1ビット目で場合分けして対処する

最後に

解説記事を書くのは初めてで不慣れな部分も多く、説明が不足していたり、わかりにくい箇所があるかもしれません。ちょっとしたことでもコメントいただけると嬉しいです！