|
| 1 | +-- https://github.com/minoki/my-atcoder-solutions |
| 2 | +{-# LANGUAGE TypeApplications #-} |
| 3 | +import Control.Monad |
| 4 | +import Data.Bits |
| 5 | +import qualified Data.ByteString.Char8 as BS |
| 6 | +import Data.Char (isSpace) |
| 7 | +import Data.Int (Int64) |
| 8 | +import Data.List (unfoldr) |
| 9 | +import Data.Ratio |
| 10 | +import qualified Test.QuickCheck as QC |
| 11 | +import Control.Exception (assert) |
| 12 | + |
| 13 | +-- comb2 n = n * (n - 1) `quot` 2 without undue overflow |
| 14 | +-- n even: comb2 n = (n `quot` 2) * (n - 1) |
| 15 | +-- n odd: comb2 n = n * (n `quot` 2) |
| 16 | +comb2 :: (Integral a, Bits a) => a -> a |
| 17 | +comb2 n = (n `shiftR` 1) * ((n - 1) .|. 1) |
| 18 | + |
| 19 | +prop_comb2 :: Integer -> QC.Property |
| 20 | +prop_comb2 n = comb2 n QC.=== n * (n - 1) `quot` 2 |
| 21 | + |
| 22 | +prop_floorSum_negate_a :: QC.NonNegative (QC.Small Int64) -> QC.Positive Int64 -> Int64 -> Int64 -> QC.Property |
| 23 | +prop_floorSum_negate_a (QC.NonNegative (QC.Small n)) (QC.Positive m) a b = |
| 24 | + let does_not_overflow = (\t -> toInteger (minBound :: Int64) <= t && t <= toInteger (maxBound :: Int64)) (toInteger b + toInteger a * (toInteger n - 1)) |
| 25 | + in does_not_overflow QC.==> floorSum n m (- a) (b + a * (n - 1)) QC.=== floorSum n m a b |
| 26 | + |
| 27 | +-- floorSum n m a b |
| 28 | +-- n: non-negative, m: positive |
| 29 | +floorSum :: Int64 -> Int64 -> Int64 -> Int64 -> Int64 |
| 30 | +floorSum n m a b | assert (n >= 0 && m > 0) False = undefined |
| 31 | +floorSum n m 0 b = n * floor (b % m) |
| 32 | +floorSum 0 m a b = 0 |
| 33 | +floorSum n 1 a b = a * comb2 n + n * b |
| 34 | +floorSum n m a b |
| 35 | + | a < 0 = floorSum n m (- a) (b + a * (n - 1)) |
| 36 | + {- |
| 37 | + | a >= m || a < 0 = case a `divMod` m of |
| 38 | + (q, a') -> q * comb2 n + floorSum n m a' b |
| 39 | +-} |
| 40 | + | let m2 = m `quot` 2 |
| 41 | + , abs a > m2 = case (a + m2) `divMod` m of |
| 42 | + (q, a') -> |
| 43 | + q * comb2 n + floorSum n m (a' - m2) b |
| 44 | + | b >= m || b < 0 = case b `divMod` m of |
| 45 | + (q, b') -> q * n + floorSum n m a b' |
| 46 | + | n > m = case n `quotRem` m of |
| 47 | + (q, n') -> (q * n - comb2 (q + 1) * m) * a + q * floorSum m m a b + floorSum n' m a b |
| 48 | + -- | n < 100 = fromInteger $ floorSum_naive n m a b |
| 49 | +-- in -- fromInteger $ floorSum_naive n m a b |
| 50 | +-- - n * t - floorSum t (- a) (- m) (- b - m) + floorSum t (- a) (- m) (b - m) |
| 51 | + | otherwise = -- 0 < a < m, 0 <= b < m, 0 < n <= m |
| 52 | + -- 0 < a < m |
| 53 | + -- sum [ fromIntegral $ length [ i | i <- [0..n-1], floor ((toInteger a * toInteger i + toInteger b) % toInteger m) >= k ] | k <- [1..(floor $ (toInteger a * (toInteger n - 1) + toInteger b) % toInteger m)] ] |
| 54 | + -- sum [ fromIntegral $ length [ i | i <- [0..n-1], i >= - floor ((- toInteger m * toInteger k + toInteger b) % toInteger a) ] | k <- [1..(floor $ (toInteger a * (toInteger n - 1) + toInteger b) % toInteger m)] ] |
| 55 | + -- sum [ n - max 0 (- floor ((- toInteger m * toInteger k + toInteger b - toInteger m) % toInteger a)) | k <- [0..(floor $ (toInteger a * (toInteger n - 1) + toInteger b) % toInteger m) - 1] ] |
| 56 | + let t = floor ((toInteger a * (toInteger n - 1) + toInteger b) % toInteger m) |
| 57 | + in n * t + floorSum t a (- m) (b - m) |
| 58 | + -- ceilSum (ceiling $ (a * (n - 1) + b) % m) a m (m - b) |
| 59 | + |
| 60 | +{- |
| 61 | +ceilSum :: Int64 -> Int64 -> Int64 -> Int64 -> Int64 |
| 62 | +ceilSum n m 0 b = n * ceiling (b % m) |
| 63 | +ceilSum 0 m a b = 0 |
| 64 | +ceilSum n 1 a b = a * (n * (n - 1) `quot` 2) + n * b |
| 65 | +ceilSum n m a b |
| 66 | + | a >= m = case a `quotRem` m of |
| 67 | + (q, a') -> q * (n * (n - 1) `quot` 2) + ceilSum n m a' b |
| 68 | + | b >= m || b < 0 = case b `divMod` m of |
| 69 | + (q, b') -> q * n + ceilSum n m a b' |
| 70 | + | n > m = case n `quotRem` m of |
| 71 | + (q, n') -> (q * n - q * (q + 1) `quot` 2 * m) * a + q * ceilSum m m a b + ceilSum n' m a b |
| 72 | + | n < 100 = fromInteger $ ceilSum_naive n m a b |
| 73 | + | otherwise = n * (n - 1) `quot` 2 - floorSum n m (m - a) (- b) -- 0 < a < m, 0 <= b < m, 0 < n <= m |
| 74 | + |
| 75 | +-} |
| 76 | +floorSum_naive :: Int64 -> Int64 -> Int64 -> Int64 -> Integer |
| 77 | +floorSum_naive n m a b = sum [ floor ((fromIntegral a * fromIntegral i + fromIntegral b) % fromIntegral m) | i <- [0..n-1] ] |
| 78 | + |
| 79 | +{- |
| 80 | +ceilSum_naive :: Int64 -> Int64 -> Int64 -> Int64 -> Integer |
| 81 | +ceilSum_naive n m a b = sum [ ceiling ((fromIntegral a * fromIntegral i + fromIntegral b) % fromIntegral m) | i <- [0..n-1] ] |
| 82 | +-} |
| 83 | +prop_floorSum :: QC.NonNegative (QC.Small Int64) -> QC.Positive Int64 -> Int64 -> Int64 -> QC.Property |
| 84 | +prop_floorSum (QC.NonNegative (QC.Small n)) (QC.Positive m) a b = QC.within (100 * 1000) $ toInteger (floorSum n m a b) QC.=== floorSum_naive n m a b |
| 85 | + |
| 86 | +prop_floorSum_r :: QC.Property |
| 87 | +prop_floorSum_r = QC.forAllShrink (QC.choose (1, 10^4)) QC.shrink $ \n -> n >= 1 QC.==> |
| 88 | + QC.forAllShrink (QC.choose (1, 10^9)) QC.shrink $ \m -> m >= 1 QC.==> |
| 89 | + QC.forAllShrink (QC.choose (0, m - 1)) QC.shrink $ \a -> |
| 90 | + QC.forAllShrink (QC.choose (0, m - 1)) QC.shrink $ \b -> |
| 91 | + QC.within (100 * 1000) $ toInteger (floorSum n m a b) QC.=== floorSum_naive n m a b |
| 92 | + |
| 93 | +{- |
| 94 | +prop_ceilSum :: QC.NonNegative (QC.Small Int64) -> QC.Positive Int64 -> Int64 -> Int64 -> QC.Property |
| 95 | +prop_ceilSum (QC.NonNegative (QC.Small n)) (QC.Positive m) a b = QC.within (100 * 1000) $ toInteger (ceilSum n m a b) QC.=== ceilSum_naive n m a b |
| 96 | +-} |
| 97 | + |
| 98 | +main = do |
| 99 | + t <- readLn @Int |
| 100 | + replicateM_ t $ do |
| 101 | + [n,m,a,b] <- map fromIntegral . unfoldr (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine |
| 102 | + print $ floorSum n m a b |
0 commit comments