Skip to content

Commit c2e8698

Browse files
committed
Fix convolution_i64
1 parent af4b04c commit c2e8698

File tree

1 file changed

+9
-8
lines changed

1 file changed

+9
-8
lines changed

src/convolution.rs

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -102,9 +102,9 @@ pub fn convolution_i64(a: &[i64], b: &[i64]) -> Vec<i64> {
102102
return vec![];
103103
}
104104

105-
let i1 = internal_math::inv_gcd(M2M3 as _, M1 as _).1;
106-
let i2 = internal_math::inv_gcd(M1M3 as _, M2 as _).1;
107-
let i3 = internal_math::inv_gcd(M1M2 as _, M3 as _).1;
105+
let (_, i1) = internal_math::inv_gcd(M2M3 as _, M1 as _);
106+
let (_, i2) = internal_math::inv_gcd(M1M3 as _, M2 as _);
107+
let (_, i3) = internal_math::inv_gcd(M1M2 as _, M3 as _);
108108

109109
let c1 = convolution_raw::<i64, M1>(a, b);
110110
let c2 = convolution_raw::<i64, M2>(a, b);
@@ -116,10 +116,11 @@ pub fn convolution_i64(a: &[i64], b: &[i64]) -> Vec<i64> {
116116
.map(|((c1, c2), c3)| {
117117
const OFFSET: &[u64] = &[0, 0, M1M2M3, 2 * M1M2M3, 3 * M1M2M3];
118118

119-
let mut x = 0;
120-
x += (c1 * i1).rem_euclid(M1 as _) * (M2M3 as i64);
121-
x += (c2 * i2).rem_euclid(M2 as _) * (M1M3 as i64);
122-
x += (c3 * i3).rem_euclid(M3 as _) * (M1M2 as i64);
119+
let mut x = [(c1, i1, M1, M2M3), (c2, i2, M2, M1M3), (c3, i3, M3, M1M2)]
120+
.iter()
121+
.map(|&(c, i, m1, m2)| c.wrapping_mul(i).rem_euclid(m1 as _).wrapping_mul(m2 as _))
122+
.fold(0, i64::wrapping_add);
123+
123124
// B = 2^63, -B <= x, r(real value) < B
124125
// (x, x - M, x - 2M, or x - 3M) = r (mod 2B)
125126
// r = c1[i] (mod MOD1)
@@ -141,7 +142,7 @@ pub fn convolution_i64(a: &[i64], b: &[i64]) -> Vec<i64> {
141142
if diff < 0 {
142143
diff += M1 as i64;
143144
}
144-
x -= OFFSET[diff.rem_euclid(5) as usize] as i64;
145+
x = x.wrapping_sub(OFFSET[diff.rem_euclid(5) as usize] as _);
145146
x
146147
})
147148
.collect()

0 commit comments

Comments
 (0)