📏

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

2022/09/12に公開

math

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

AtCoder Regular Contest 148 解説記事:

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

この記事で扱う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

概要:

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^2$ 、 $R\lt N$ であるため $T/R\ge N$ があり得ます)。 (9-14行目)

最終処理では、まずは2回目の $\operatorname{iMR}$ の最後の処理で $/R$ する直前の値に対して $\mod R$ を行う事で、 $t\lt 0$ であれば $Q=2^{64}\mod R=932051910$ 、 $t\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$ 以上か否か)を考慮しないといけません。また、 $R$ が $2^{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