Official

F - 998244353 → 1000000007 Editorial by Nyaan


この問題は基本的にモンゴメリ乗算のアルゴリズムを利用すれば解くことができます。モンゴメリ乗算の詳細な説明は Wikipedia などの記事に譲ります。(以下の解説では、変数への文字の割り当て方は Wikipedia に準じます)

モンゴメリ乗算を端的に説明すると、 \(A \times B \bmod{N}\) を乗算・加算・\(\text{mod }R\) の剰余算のみを用いて(\(\text{mod }N\) の剰余算を利用せずに)計算するアルゴリズムです。次のコードは \(C\)\(A \times B \bmod{N}\) を代入しています。
(以下では \(N'\)\(N' N \equiv -1 \pmod{R}\) を満たす最小の正整数とします。また、\(R_2 = R^2 \bmod{N}\) とします。)

uint64_t reduce(uint64_t T) {
  uint64_t t = (T + (T % R * Nprime % R) * N) / R;
  if (t >= N) t -= N;
  return t;
}

uint64_t A = a, B = b, C;
void f() {
  A = reduce(A * R_2);
  B = reduce(B * R_2);
  C = A * B;
  C = reduce(C);
  C = reduce(C);
}

\(N = 10^9 + 7, R = 998244353\) としてモンゴメリ乗算を行うことを考えます。

  • 追記:この部分は少し軽率な展開でした。montgomery reduction では 入力の \(t\)\(0 \leq t \lt NR\) であるのを仮定していますが、この問題では \(N \gt R\) なので 入力の \(t\)\(NR\) 以上になる場合があります。ただし、\(R_2 \lt N/3\) 等を利用して乗算の過程における誤差を正確に評価することで、上のコードが正しく動くことが証明できます。

上のコードの reduce(T) の部分は montgomery reduction と呼ばれています。montgomery reduction をうまく計算できればこの問題を解くことができます。コードを読むと、次の 2 か所の部分を除けばすべて所定の命令を用いて実装できることがわかります。

  • \(1\) 行目の / R の部分
  • \(2\) 行目の if

1 つ目の \(R\) で割る部分については、乗算が \(\text{mod }2^{64}\) を取ることを利用すれば比較的容易に解決します。
\(S\)\(S R \equiv 1 \pmod{2^{64}}\) を満たす整数とします。\(R\) の倍数であるような任意の 64 bit 整数 \(X\) について \(X / R\)\(X S\)\(\text{mod }2^{64}\) 上で合同で、かつそのような 値は \(0\) 以上 \(2^{64}\) 未満で一意に定まります。よって \(R\) で割る代わりに \(S\) を掛けても同じ値を得られることが示されます。

次に if 文の部分をうまく処理する方法を考えます。これはいくつかの方法が考えられます。

解法 1

1 つ目の解法は正攻法で、if 文をうまく実装します。問題を単純にするために以下の関数 \(F(t)\) を考えます。

\[F(t) = (1 \text{ if t }\geq N\text{ else }0)\]

\(F(t)\) を所定の演算のみを用いて計算できれば、if 文の箇所は \(t + F(t) \times (2^{64} - 1)\) に置き換えられます。

\(F(t)\) をうまく実装します。これは \(0 \leq t \lt N\) の場合と \(N \leq t \lt 2N\) の場合で異なる結果が得られる計算を行えばよく、例えば次のような方法が考えられます。

以下の \(p, q\) を変数に格納することを考えます。

\[p = t + 2^{64} - N, q = t + 2^{64} + 2R - N\]

\(t \lt N\) の場合は \(q\) のみが \(2^{64}\) 以上になるので、\(q\) に格納される値は真の値から \(2^{64}\) を引いた値になります。よって \((q - p) \bmod{R}\)\(2^{64} \bmod{R}\) になります。
一方 \(N \leq t \lt 2N\) の場合は \(p, q\) ともに \(2^{64}\) 以上になり、\(2^{64}\) を引かれます。よって \((q - p) \bmod{R}\)\(0\) です。
よって、 \((q-p) \bmod{R}\) に適切に定数を乗算・加算すれば \(F(t)\) を計算できることがわかりました。(\(2^{64} \bmod{R}\)\(\text{mod }2^{64}\) 上で逆元を持たないことから、この部分は見た目より少し複雑な実装が要りますがここでは省略します。)これを利用すれば reduce 関数を完全に実装できて、 reduce 関数を用いてモンゴメリ乗算を行えば当初の目的を達成できます。

解法 2

もう 1 つの解法 (当初の想定解) は力技です。
モンゴメリ乗算では \(T R^{-1}\)\(0\) 以上 \(N\) 未満の範囲の値で管理しています。montgomery reduction での if 文は、1 文前の \(t\)\(0\) 以上 \(2N\) 未満の値を取るので、これを \(\lbrack 0, N)\) の範囲に正規化するために使用されているものです。
逆に言うと、if 文を取らなくても (オーバーフロー等が発生しなければ) 常に真の値と \(\text{mod }N\) において一致する答えを得られる、ということも言えます。そこで、reduce において正規化を行わずに乗算を行ってみましょう。すなわち、次のような incomplete_reduce 関数で表せる操作を考えます。

uint64_t incomplete_reduce(uint64_t T) {
  return (T + (T % R * Nprime % R) * N) / R;
}

reduce 関数を incomplete_reduce に置き換えたコードを考えます。(以下に再掲します。) f\(5\) 行目の \(C\)\(N\) 未満に収まっていれば、上の操作をプログラムに直したものが正答となります。(実際はなりません)

uint64_t A = a, B = b, C;
void f() {
  A = incomplete_reduce(A * R_2);
  B = incomplete_reduce(B * R_2);
  C = A * B;
  C = incomplete_reduce(C);
  C = incomplete_reduce(C);
}

各行について、左辺に代入される値を評価していくと、詳細は略しますが、5 行目の C\(N+6\) 以下になるのが確認できます。(incomplete_reduce の返り値が \((T + (R-1)N) / R\) で抑えられるのを利用すれば示せます。なお、実際はもう少し厳しく評価できます。)

ところで、incomplete_reduce の返り値は \(T\) に依存するので、入力の \(T\) が小さいと出力の値もより厳しく評価できるのがわかります。そこで reduction を もう 1 回行います。すなわち次の処理を \(6\) 行目, \(7\) 行目に加えます。

C = incomplete_reduce(C * R_2);
C = incomplete_reduce(C);

このような処理を加えると、これも詳細は略しますが、\(C \lt N+1\) になるのが確認できます。(\(R_2 \lt N/3\) なのを利用する必要があると思っています。) また、\(C=N\) になる可能性はあり得ません。(その場合 \(a,b\) の少なくとも一方が \(0\) だが、incomplete_reduce(0) = 0 なので \(C=0\) になる。) よって \(C\)\(N\) 未満に収まります。
よって、不完全な reduction を 1 回余分に行えば \(C\)\(N\) 未満に収まり、\(a \times b \bmod{N}\) を計算できるということが証明できました。

以上より、解法 1, 解法 2 のいずれかをプログラムに書けば、\(a \times b \pmod{N}\) を計算できます。どちらの解法も行数に十分余裕があります。(なお、準備陣の行数は 27 行, 37 行, 40 行, 40 行でした) ちなみに、はじめに \(A \times B\) を計算して、その値に対して reduction を行うことで命令数を削減できます。

なお、解法 1 の説明は maspy さんのプログラムを, 命令数を減らせる点については maroon さんのプログラムを参考にしました。admin, tester の皆様ありがとうございました。

posted:
last update: