From 38624c9f991f38ab8b9e1e5ad477648fe4707d4a Mon Sep 17 00:00:00 2001 From: nebocco Date: Thu, 9 Jun 2022 14:54:03 +0900 Subject: [PATCH 1/6] fix: relax the constraints of floor_sum --- src/internal_math.rs | 35 +++++++++++++++++++++++++++++++++++ src/math.rs | 44 +++++++++++++++++++++++++++++--------------- 2 files changed, 64 insertions(+), 15 deletions(-) diff --git a/src/internal_math.rs b/src/internal_math.rs index 515191c..3742544 100644 --- a/src/internal_math.rs +++ b/src/internal_math.rs @@ -235,6 +235,41 @@ pub(crate) fn primitive_root(m: i32) -> i32 { // omitted // template constexpr int primitive_root = primitive_root_constexpr(m); +/// # Arguments +/// * `n` `n < 2^32` +/// * `m` `1 <= m < 2^32` +/// +/// # Returns +/// `sum_{i=0}^{n-1} floor((ai + b) / m) (mod 2^64)` +/* const */ +#[allow(clippy::many_single_char_names)] +pub(crate) fn floor_sum_unsigned(mut n: u64, mut m: u64, mut a: u64, mut b: u64) -> u64 { + let mut ans = 0; + loop { + if a >= m { + if n > 0 { + ans += n * (n - 1) / 2 * (a / m); + } + a %= m; + } + if b >= m { + ans += n * (b / m); + b %= m; + } + + let y_max = a * n + b; + if y_max < m { + break; + } + // y_max < m * (n + 1) + // floor(y_max / m) <= n + n = y_max / m; + b = y_max % m; + std::mem::swap(&mut m, &mut a); + } + return ans; +} + #[cfg(test)] mod tests { #![allow(clippy::unreadable_literal)] diff --git a/src/math.rs b/src/math.rs index 61f15d5..3364a12 100644 --- a/src/math.rs +++ b/src/math.rs @@ -186,24 +186,20 @@ pub fn crt(r: &[i64], m: &[i64]) -> (i64, i64) { /// assert_eq!(math::floor_sum(6, 5, 4, 3), 13); /// ``` pub fn floor_sum(n: i64, m: i64, mut a: i64, mut b: i64) -> i64 { + assert!(0 <= n && n < 1i64 << 32); + assert!(1 <= m && m < 1i64 << 32); let mut ans = 0; - if a >= m { - ans += (n - 1) * n * (a / m) / 2; - a %= m; + if a < 0 { + let a2 = internal_math::safe_mod(a, m); + ans -= n * (n - 1) / 2 * ((a2 - a) / m); + a = a2; } - if b >= m { - ans += n * (b / m); - b %= m; + if b < 0 { + let b2 = internal_math::safe_mod(b, m); + ans -= n * ((b2 - b) / m); + b = b2; } - - let y_max = (a * n + b) / m; - let x_max = y_max * m - b; - if y_max == 0 { - return ans; - } - ans += (n - (x_max + a - 1) / a) * y_max; - ans += floor_sum(y_max, a, m, (a - x_max % a) % a); - ans + ans + internal_math::floor_sum_unsigned(n as u64, m as u64, a as u64, b as u64) as i64 } #[cfg(test)] @@ -306,5 +302,23 @@ mod tests { 499_999_999_500_000_000 ); assert_eq!(floor_sum(332955, 5590132, 2231, 999423), 22014575); + for n in 0..20 { + for m in 1..20 { + for a in -20..20 { + for b in -20..20 { + assert_eq!(floor_sum(n, m, a, b), floor_sum_naive(n, m, a, b)); + } + } + } + } + } + + fn floor_sum_naive(n: i64, m: i64, a: i64, b: i64) -> i64 { + let mut ans = 0; + for i in 0..n { + let z = a * i + b; + ans += (z - internal_math::safe_mod(z, m)) / m; + } + ans } } From 4e6065b74a5216063638a65c74641b389ba66090 Mon Sep 17 00:00:00 2001 From: nebocco Date: Fri, 10 Jun 2022 17:29:26 +0900 Subject: [PATCH 2/6] fmt: cargo clippy --- src/internal_math.rs | 2 +- src/math.rs | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/internal_math.rs b/src/internal_math.rs index 3742544..cb76908 100644 --- a/src/internal_math.rs +++ b/src/internal_math.rs @@ -267,7 +267,7 @@ pub(crate) fn floor_sum_unsigned(mut n: u64, mut m: u64, mut a: u64, mut b: u64) b = y_max % m; std::mem::swap(&mut m, &mut a); } - return ans; + ans } #[cfg(test)] diff --git a/src/math.rs b/src/math.rs index 3364a12..e41705f 100644 --- a/src/math.rs +++ b/src/math.rs @@ -185,6 +185,7 @@ pub fn crt(r: &[i64], m: &[i64]) -> (i64, i64) { /// /// assert_eq!(math::floor_sum(6, 5, 4, 3), 13); /// ``` +#[allow(clippy::many_single_char_names)] pub fn floor_sum(n: i64, m: i64, mut a: i64, mut b: i64) -> i64 { assert!(0 <= n && n < 1i64 << 32); assert!(1 <= m && m < 1i64 << 32); @@ -313,6 +314,7 @@ mod tests { } } + #[allow(clippy::many_single_char_names)] fn floor_sum_naive(n: i64, m: i64, a: i64, b: i64) -> i64 { let mut ans = 0; for i in 0..n { From ac9ab7e99ee33c81b9631703f67ab2b8c793a361 Mon Sep 17 00:00:00 2001 From: nebocco Date: Mon, 13 Jun 2022 19:24:48 +0900 Subject: [PATCH 3/6] fmt: clippy --- src/math.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/math.rs b/src/math.rs index e41705f..7f4bd1f 100644 --- a/src/math.rs +++ b/src/math.rs @@ -187,8 +187,8 @@ pub fn crt(r: &[i64], m: &[i64]) -> (i64, i64) { /// ``` #[allow(clippy::many_single_char_names)] pub fn floor_sum(n: i64, m: i64, mut a: i64, mut b: i64) -> i64 { - assert!(0 <= n && n < 1i64 << 32); - assert!(1 <= m && m < 1i64 << 32); + assert!((0..1i64 << 32).contains(&n)); + assert!((1..1i64 << 32).contains(&m)); let mut ans = 0; if a < 0 { let a2 = internal_math::safe_mod(a, m); From 7703e17d5107b37d030594b716099e5c8c35693e Mon Sep 17 00:00:00 2001 From: TonalidadeHidrica <47710717+TonalidadeHidrica@users.noreply.github.com> Date: Sat, 15 Apr 2023 21:08:47 +0200 Subject: [PATCH 4/6] Update docs of floor_sum --- src/math.rs | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/math.rs b/src/math.rs index 7f4bd1f..bd92feb 100644 --- a/src/math.rs +++ b/src/math.rs @@ -162,13 +162,16 @@ pub fn crt(r: &[i64], m: &[i64]) -> (i64, i64) { (r0, m0) } -/// Returns $\sum_{i = 0}^{n - 1} \lfloor \frac{a \times i + b}{m} \rfloor$. +/// Returns +/// +/// $$\sum_{i = 0}^{n - 1} \left\lfloor \frac{a \times i + b}{m} \right\rfloor.$$ +/// +/// It returns the answer in $\bmod 2^{\mathrm{64}}$, if overflowed. /// /// # Constraints /// -/// - $0 \leq n \leq 10^9$ -/// - $1 \leq m \leq 10^9$ -/// - $0 \leq a, b \leq m$ +/// - $0 \leq n \lt 2^{32}$ +/// - $1 \leq m \lt 2^{32}$ /// /// # Panics /// @@ -176,7 +179,7 @@ pub fn crt(r: &[i64], m: &[i64]) -> (i64, i64) { /// /// # Complexity /// -/// - $O(\log(n + m + a + b))$ +/// - $O(\log{(m+a)})$ /// /// # Example /// From ece6abe22f683720999d527b42c4f0fd435e7794 Mon Sep 17 00:00:00 2001 From: TonalidadeHidrica <47710717+TonalidadeHidrica@users.noreply.github.com> Date: Sat, 15 Apr 2023 22:32:28 +0200 Subject: [PATCH 5/6] In floor_sum, use Wrapping to handle overflows --- src/internal_math.rs | 15 ++++++++++----- src/math.rs | 21 ++++++++++++--------- 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/src/internal_math.rs b/src/internal_math.rs index cb76908..c6290b0 100644 --- a/src/internal_math.rs +++ b/src/internal_math.rs @@ -1,6 +1,6 @@ // remove this after dependencies has been added #![allow(dead_code)] -use std::mem::swap; +use std::{mem::swap, num::Wrapping as W}; /// # Arguments /// * `m` `1 <= m` @@ -243,12 +243,17 @@ pub(crate) fn primitive_root(m: i32) -> i32 { /// `sum_{i=0}^{n-1} floor((ai + b) / m) (mod 2^64)` /* const */ #[allow(clippy::many_single_char_names)] -pub(crate) fn floor_sum_unsigned(mut n: u64, mut m: u64, mut a: u64, mut b: u64) -> u64 { - let mut ans = 0; +pub(crate) fn floor_sum_unsigned( + mut n: W, + mut m: W, + mut a: W, + mut b: W, +) -> W { + let mut ans = W(0); loop { if a >= m { - if n > 0 { - ans += n * (n - 1) / 2 * (a / m); + if n > W(0) { + ans += n * (n - W(1)) / W(2) * (a / m); } a %= m; } diff --git a/src/math.rs b/src/math.rs index bd92feb..d5c58df 100644 --- a/src/math.rs +++ b/src/math.rs @@ -189,21 +189,24 @@ pub fn crt(r: &[i64], m: &[i64]) -> (i64, i64) { /// assert_eq!(math::floor_sum(6, 5, 4, 3), 13); /// ``` #[allow(clippy::many_single_char_names)] -pub fn floor_sum(n: i64, m: i64, mut a: i64, mut b: i64) -> i64 { +pub fn floor_sum(n: i64, m: i64, a: i64, b: i64) -> i64 { + use std::num::Wrapping as W; assert!((0..1i64 << 32).contains(&n)); assert!((1..1i64 << 32).contains(&m)); - let mut ans = 0; + let mut ans = W(0 as u64); + let (wn, wm, mut wa, mut wb) = (W(n as u64), W(m as u64), W(a as u64), W(b as u64)); if a < 0 { - let a2 = internal_math::safe_mod(a, m); - ans -= n * (n - 1) / 2 * ((a2 - a) / m); - a = a2; + let a2 = W(internal_math::safe_mod(a, m) as u64); + ans -= wn * (wn - W(1)) / W(2) * ((a2 - wa) / wm); + wa = a2; } if b < 0 { - let b2 = internal_math::safe_mod(b, m); - ans -= n * ((b2 - b) / m); - b = b2; + let b2 = W(internal_math::safe_mod(b, m) as u64); + ans -= wn * ((b2 - wb) / wm); + wb = b2; } - ans + internal_math::floor_sum_unsigned(n as u64, m as u64, a as u64, b as u64) as i64 + let ret = ans + internal_math::floor_sum_unsigned(wn, wm, wa, wb); + ret.0 as i64 } #[cfg(test)] From 46caad4bc6f7f01881cf7b405719e167dee27dba Mon Sep 17 00:00:00 2001 From: TonalidadeHidrica <47710717+TonalidadeHidrica@users.noreply.github.com> Date: Sat, 15 Apr 2023 23:27:45 +0200 Subject: [PATCH 6/6] Fix clippy --- src/math.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/math.rs b/src/math.rs index d5c58df..5dfd349 100644 --- a/src/math.rs +++ b/src/math.rs @@ -193,7 +193,7 @@ pub fn floor_sum(n: i64, m: i64, a: i64, b: i64) -> i64 { use std::num::Wrapping as W; assert!((0..1i64 << 32).contains(&n)); assert!((1..1i64 << 32).contains(&m)); - let mut ans = W(0 as u64); + let mut ans = W(0_u64); let (wn, wm, mut wa, mut wb) = (W(n as u64), W(m as u64), W(a as u64), W(b as u64)); if a < 0 { let a2 = W(internal_math::safe_mod(a, m) as u64);