@@ -20,26 +20,19 @@ comb2 n = (n `shiftR` 1) * ((n - 1) .|. 1)
2020-- floorSum n m a b
2121-- Assumptions:
2222-- * n: non-negative, m: positive
23+ -- * a and b can be negative, or >= m
2324floorSum :: Int64 -> Int64 -> Int64 -> Int64 -> Int64
2425floorSum ! n ! m ! a ! b
2526 | assert (n >= 0 && m > 0 ) False = undefined
26- | a < 0 = floorSum n m (- a) (b + a * (n - 1 ))
27- {-
28- | m < n = case n `quotRem` m of
29- (q, n') -> (q * n - comb2 (q + 1) * m) * a +
30- q * floorSum_positive 0 m m a b +
31- floorSum_positive 0 n' m a b
32- -}
27+ | a < 0 = floorSum_positive 0 n m (- a) (b + a * (n - 1 ))
3328 | otherwise = floorSum_positive 0 n m a b
3429 where
3530 -- Invariants:
3631 -- * n: non-negative, m: positive, a: non-negative
3732 -- * 0 <= n <= m
3833 floorSum_positive :: Int64 -> Int64 -> Int64 -> Int64 -> Int64 -> Int64
3934 floorSum_positive ! acc ! n ! m ! a ! b
40- | a == 0 = acc + n * (b `div` m)
4135 | n == 0 = acc
42- | m == 1 = acc + a * comb2 n + b * n
4336 | let m2 = m `quot` 2 , a > m2 =
4437 let (q, a') = (a + m2) `quotRem` m
4538 (a'',b'') = if a' < m2 then
0 commit comments