📏

AtCoder Regular Contest 148 - F - 998244353 → 1000000007 解法例

2022/09/12に公開約6,300字

AtCoder Regular Contest 148 - F 「998244353 → 1000000007」 問題本文:

https://atcoder.jp/contests/arc148/tasks/arc148_f

AtCoder Regular Contest 148 解説記事:

https://atcoder.jp/contests/arc148/editorial

snukeさんによる18行解法(現在は17行解法が追加されています):

https://atcoder.jp/contests/arc148/editorial/4831

この記事で扱う18行解法の変種:

18
mul T A B
rem U T
mul U U 4915446
rem U U
mul U U 1000000007
add U U T
mul T U 8456992263029997534
rem U T
mul U U 993328907
rem U U
mul U U 18446744072709551609
add U U T
mul T U 996491785301655553
rem U U
mul V U 696235320
rem V V
add U U V
add C T U

https://atcoder.jp/contests/arc148/submissions/34818980

概要:

1回目に if (t >= N) t-= N; の無い incomplete montgomery reduction (\operatorname{iMR}) を行い、
2回目はそのような if文 を用いる montgomery reduction (\operatorname{MR}) を行い、
答えを \operatorname{MR}(\operatorname{iMR}(A*B)*R_2) として求める方針はユーザ解説の snukeさんの18行解法と変わりません。

但し、2回目の montgomery reduction では、 0\le T より得られる値 t の範囲が \left\{-N\lt t\le(T/R)\right\} となる \operatorname{iMR} の変種 (if (t >= N) t-= N; の代わりに if (t < 0) t += N; のような後処理になる) を用いている点が異なります。-N\lt t\lt 0 の範囲の数は、変数上では 2^{64}-N\lt t\lt 2^{64} に存在します。また、t の上界についてもこの解法の2回目のモンゴメリリダクションの範囲においては T\lt (((N^2/R) + N)*R_2)\lt NR であるため、 T/R\lt N を満たしています (1回目は 0\le T\lt N^2R\lt N であるため T/R\ge N があり得ます)。 (9-14行目)

最終処理では、まずは2回目の \operatorname{iMR} の最後の処理で /R する直前の値に対して \mod R を行う事で、 t\lt 0 であれば Q=2^{64}\mod R=932051910t\ge 0 であれば 0 の値が抽出されます。 u\leftarrow\text{if }t\lt 0\text{ then }Q,\text{ else }0 (15行目)

残念ながら Q=2^{64}\mod R=932051910 の値と 2^{64} は互いに素ではないため、これを N の値に変換するためには工数が掛かります。
u=\text{if }t\lt 0\text{ then }Q,\text{ else }0 に対し、 u\leftarrow u + ((u*696235320)\mod R) の処理を行うことで u=\text{if }t\lt 0\text{ then }N,\text{ else }0 の状態に変換します。(16-18行目)

最後に、 c\leftarrow t+u とする事で最終的な答えとしています。(19行目)

1. 18
2. mul T A B T\leftarrow a*b
3. rem U T 3-8行目: t\leftarrow (T+(((T\mod R)*N'\mod R) * N))/R\\\left\{T\ge 0,\ 0\le t\lt\left((T/R)+N\right)\right\}
4. mul U U 4915446
5. rem U U
6. mul U U 1000000007
7. add U U T
8. mul T U 8456992263029997534 T\leftarrow t*R_2 (リダクション1回目の /R 演算部分と融合)
9. rem U T 9-14行目: t\leftarrow (T-(((T\mod R)*N^{-1}\mod R) * N))/R\\\left\{T\ge 0,\ -N\lt t\le(T/R)\right\}
10. mul U U 993328907
11. rem U U
12. mul U U 18446744072709551609
13. add U U T
14. mul T U 996491785301655553
15. rem U U 15-18行目: u\leftarrow\text{if }t\lt 0\text{ then }N,\text{ else }0
16. mul V U 696235320
17. rem V V
18. add U U V
19. add C T U c\leftarrow\text{if }t\lt 0\text{ then }(t+N),\text{ else }(t+0)

定数・マジックナンバー:

  • N=1000000007
  • R=998244353
  • -N\mod 2^{64}=18446744072709551609
  • N'=R-(N^{-1}\mod R)=4915446
  • N^{-1}\mod R=993328907
  • R^{-1}\mod 2^{64}=996491785301655553
  • R_2=R^2\mod N=320946142
  • R_c=R_2*R^{-1}\mod 2^{64}=8456992263029997534
  • Q=2^{64}\mod R=932051910
  • P=(N-Q)/Q\mod R=696235320

概念コード:

/*
18
mul T A B
rem U T
mul U U 4915446
rem U U
mul U U 1000000007
add U U T
mul T U 8456992263029997534
rem U T
mul U U 993328907
rem U U
mul U U 18446744072709551609
add U U T
mul T U 996491785301655553
rem U U
mul V U 696235320
rem V V
add U U V
add C T U
*/

const N: u64 = 1000000007;
const R: u64 = 998244353;
const NEG_N: u64 = 18446744072709551609; // 2 ** 64 - n; 
const N_DASH: u64 = 4915446; // n * ndash % r == r - 1
const N_INV: u64 = 993328907; // n * ninv % r == 1
const R_INV: u64 = 996491785301655553; // r * rinv % (2**64) == 1
const R2: u64 = 320946142; // == r * r % n
//const RC: u64 = 8456992263029997534; // R_INV * R2 % (2**64)
const Q: u64 = 932051910; // 2**64 % R
const P: u64 = 696235320; // Q * P % R + Q == N

fn reduction1(t: u64) -> u64 {
    // t := (T + (((T % R) * N_DASH % R) * N)) * R_INV (0 <= t < T/R + N)
    let u = t % R;
    let u = u.wrapping_mul(N_DASH);
    let u = u % R;
    let u = u.wrapping_mul(N);
    let u = u.wrapping_add(t);
    let _t = u.wrapping_mul(R_INV);
    assert_eq!(_t * R % N, t % N);
    _t
}

fn reduction2(t: u64) -> (u64, u64) { // t の値域が異なるモンゴメリ剰余乗算の変種を使う
    // t := (T + (((T % R) * N_INV % R) * NEG_N)) * R_INV (-N < t <= T/R)
    let u = t % R;
    let u = u.wrapping_mul(N_INV);
    let u = u % R;
    let u = u.wrapping_mul(NEG_N);
    let u = u.wrapping_add(t);
    let _t = u.wrapping_mul(R_INV);
    assert_eq!(_t.wrapping_add(N) * R % N, t % N);
    (_t, u)
}

fn fix(t: u64, u: u64) -> u64 { // if t < 0 then return t + N else return t + 0
    //let u = t.wrapping_mul(R); // reduction2 の最後で /R しているのでその前から値を貰う
    let u = u % R; // t が負なら Q=2**64%R 、 t が非負なら 0 の値を u に代入
    assert!(((t as i64) < 0 && u == Q) || ((t as i64) >= 0 && u == 0));
    let v = u.wrapping_mul(P);
    let v = v % R;
    let u = u.wrapping_add(v); // t が負なら N 、 t が非負なら 0 の値を u に代入
    assert!(((t as i64) < 0 && u == N) || ((t as i64) >= 0 && u == 0));
    let _c = t.wrapping_add(u); // t が負なら t+N 、 t が非負なら t+0 の値を _c に代入
    _c
}

pub fn mulmod(a: u64, b: u64) -> u64 {
    let t = reduction1(a * b);
    let (t, u) = reduction2(t * R2);
    let c = fix(t, u);
    assert_eq!(a * b % N, c);
    c
}

モンゴメリ剰余乗算でこの2回目に用いたような変種を用いると、 if (t >= N) t-= N;if (t < 0) t += N; のように置き換えられ、さらに演算オーバーフローの検出と組み合わせた処理に出来るケースがあるかと思いますので、ちょっと覚えてもらえればと。例えば、法が64bit長整数での剰余乗算を行う時はこの変種を使わないと、64bit整数を上にオーバーフローするケースの場合分け(64bit境界を上にオーバーフローしているか、オーバーフローしていないが N 以上か否か)を考慮しないといけません。また、 R2^{32}2^{64} の累乗数などを使うことで効率化可能な状況(例えば多倍長整数型だとリアロケーションやワード移動処理が絡んでしまい当てはまらないケースがあります)の場合などは更に床関数 \lfloor x\rfloor を用いて加減算の前にそれぞれ除算してしまうことで
t \leftarrow \lfloor{T/R}\rfloor - \lfloor{(TN^{-1} \mod R)N/R}\rfloor;
\text{if }(t\lt 0)\text{ then \{ return }(t+N)\text{ \} else \{ return }(t)\text{ \}};
と変形することで下位ビットの処理を簡略化することができます(実際は除算+床関数の部分を右ビットシフト処理をするようにコーディング、もしくはワード境界の処理範囲を調整)。

さらに余談ですが、gcc/clang の場合は、 __builtin_uaddll_overflow,__builtin_usubll_overflow,__builtin_umulll_overflow などのオーバーフローチェック付き演算のビルトイン関数が存在します。Rust言語 の場合は、整数プリミティブに overflowing_add,overflowing_sub,overflowing_mul などの実装が存在します。これらのオーバーフローチェック付き演算を併用することで、{キャリーフラグ/オーバーフローフラグ/...}チェック付き{ジャンプ/コピー}命令、などのハードウェア命令を併用した処理にコンパイルされやすくなります。

Discussion

ログインするとコメントできます