-- | Extra math module.
--
-- @since 1.0.0.0
module AtCoder.Extra.Math
  ( -- * Re-exports from the internal math module
    isPrime32,
    ACIM.invGcd,
    primitiveRoot32,

    -- * Prime numbers and divisors
    primes,
    isPrime,
    primeFactors,
    primeFactorsUnsorted,
    divisors,
    divisorsUnsorted,

    -- * Binary exponentiation

    -- | ==== __Examples__
    -- >>> import AtCoder.Extra.Math qualified as M
    -- >>> import Data.Semigroup (Product(..), Sum(..))
    -- >>> getProduct $ M.power (<>) 32 (Product 2)
    -- 4294967296
    --
    -- >>> getProduct $ M.stimes' 32 (Product 2)
    -- 4294967296
    --
    -- >>> getProduct $ M.mtimes' 32 (Product 2)
    -- 4294967296
    power,
    stimes',
    mtimes',
  )
where

import AtCoder.Extra.Math.Montgomery64 qualified as M64
import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.Internal.Math qualified as ACIM
import Control.Monad (unless, when)
import Data.Bit (Bit (..))
import Data.Bits (bit, countTrailingZeros, (.<<.), (.>>.))
import Data.Foldable (for_)
import Data.Maybe (fromJust)
import Data.Vector.Algorithms.Intro qualified as VAI
import Data.Vector.Algorithms.Radix qualified as VAR
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import Data.Word (Word64)
import GHC.Stack (HasCallStack)
import System.Random

-- | \(O(k \log^3 n) (k = 3)\). Returns whether the given `Int` value is a prime number.
--
-- ==== Constraints
-- - \(n < 4759123141 (2^{32} < 4759123141)\), otherwise the return value can lie
--   (see [Wikipedia](https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test#Testing_against_small_sets_of_bases)).
--
--
-- @since 1.1.0.0
{-# INLINE isPrime32 #-}
isPrime32 :: (HasCallStack) => Int -> Bool
isPrime32 :: HasCallStack => Int -> Bool
isPrime32 Int
x = Int -> Bool
ACIM.isPrime Int
x
  where
    !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
4759123141) (String -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.Math.isPrime32: given too large number `" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
x String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"`"

-- | Returns the primitive root of the given `Int`.
--
-- ==== Constraints
-- - The input must be a prime number.
-- - The input must be less than \(2^32\).
--
-- @since 1.2.0.0
{-# INLINE primitiveRoot32 #-}
primitiveRoot32 :: (HasCallStack) => Int -> Int
primitiveRoot32 :: HasCallStack => Int -> Int
primitiveRoot32 Int
x = Int -> Int
ACIM.primitiveRoot Int
x
  where
    !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
32)) (String -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.Math.primitiveRoot32: given too large number `" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
x String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"`"

-- | \(O(n \log \log n)\) Creates an array of prime numbers up to the given limit, using Sieve of
-- Eratosthenes.
--
-- The minimum computational complexity is \(\Omega(B \log \log B)\), where \(B = 2^{15}\) is the
-- length of segment. This constraint comes from the use of segmented sieve.
--
-- ==== Constraints
-- - The upper limit must be less than or equal to \(2^{30} (\gt 10^9)\), otherwise the returned
-- prime table is incorrect.
--
-- @since 1.2.6.0
{-# INLINEABLE primes #-}
primes :: Int -> VU.Vector Int
primes :: Int -> Vector Int
primes Int
upperLimit
  | Int
upperLimit Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = Vector Int
forall a. Unbox a => Vector a
VU.empty
  | Bool
otherwise = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
      -- segment length (TODO: isn't it 32767?)
      let !s :: Int
s = Int
32768 :: Int -- 2 ^ 15

      -- sieve length (TODO: use \sqrt limit? do benchmark)
      let !sieveMax :: Int
sieveMax = Int
s

      -- Is it like LT bound??
      let !limit :: Int
limit = Int
upperLimit Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

      -- base primes with index
      (!Vector Int
ps, !MVector (PrimState (ST s)) Int
is) <- do
        MVector s Bit
sieve <- Int -> Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
sieveMax Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Bit -> ST s (MVector (PrimState (ST s)) Bit))
-> Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
False
        MVector s Int
ps <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int
sieveMax Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
        MVector (PrimState (ST s)) Int
is <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int
sieveMax Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
        -- FIXME: carry index?
        MVector s Int
iNext <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 (Int
0 :: Int)
        [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
3, Int
5 .. Int
s] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
p1 -> do
          Bit Bool
b <- MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Bit
MVector (PrimState (ST s)) Bit
sieve Int
p1
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
b (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            Int
at <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
iNext Int
0
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
iNext Int
0 (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
at Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            -- (base prime, next index (odd numbers only, so `div` 2)
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
ps Int
at Int
p1
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
is Int
at (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
p1 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p1 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
            -- NOTE: if `j` is a composite number, it's already enumerated by a smaller prime
            -- number than `p0`, so skip to `p1 * p1` and iterate through odd numbers only:
            [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
p1 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p1, Int
p1 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p1 .. Int
sieveMax] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
np1 -> do
              MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Bit
MVector (PrimState (ST s)) Bit
sieve Int
np1 (Bit -> ST s ()) -> Bit -> ST s ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
        Int
len <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
iNext Int
0
        (,Int
-> MVector (PrimState (ST s)) Int -> MVector (PrimState (ST s)) Int
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VGM.take Int
len MVector (PrimState (ST s)) Int
is) (Vector Int -> (Vector Int, MVector (PrimState (ST s)) Int))
-> ST s (Vector Int)
-> ST s (Vector Int, MVector (PrimState (ST s)) Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (Int -> MVector s Int -> MVector s Int
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VGM.take Int
len MVector s Int
ps)

      -- https://en.wikipedia.org/wiki/Prime-counting_function
      let !Int
maxPrimeCount :: Int
            -- NOTE: 1,700 is a point where the next function estimates better as far as I tested:
            | Int
limit Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1700 = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double
1.25506 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
limit Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
log (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
limit) :: Double)
            -- Rosser and Schoenfeld Boundsh (1962): holds for x > e^{3/2}:
            | Int
limit Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
60184 = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
limit Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double -> Double
forall a. Floating a => a -> a
log (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
limit) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1.5) :: Double)
            -- Pierre Dusart (2010): holds for x >= 60184:
            | Bool
otherwise = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
limit Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double -> Double
forall a. Floating a => a -> a
log (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
limit) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1.1) :: Double)

      -- let f x = round (1.25506 * fromIntegral x / log (fromIntegral x) :: Double)
      -- let g x = round (fromIntegral x / (log (fromIntegral x) - 1.5) :: Double)
      -- let h x = ceiling (fromIntegral x / (log (fromIntegral x) - 1.1) :: Double)
      -- let p x = (f x, g x, h x)

      MVector s Int
result <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
maxPrimeCount (-Int
1)
      MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
result Int
0 Int
2
      -- FIXME: carry index?
      MVector s Int
nPrimes <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 (Int
1 :: Int)

      -- Sieve of Eratosthenes by block of size `s`, ignoring even numers
      -- FIXME: block length of size `s/2` should make more sense?
      MVector (PrimState (ST s)) Bit
block <- Int -> ST s (MVector (PrimState (ST s)) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
s
      let !r :: Int
r = Int
limit Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
      [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
1, Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s .. Int
r] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
l -> do
        MVector (PrimState (ST s)) Bit -> Bit -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
VGM.set MVector (PrimState (ST s)) Bit
block (Bit -> ST s ()) -> Bit -> ST s ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
False

        Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Int
ps ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
idx Int
p -> do
          -- FIXME: cut out the ps beforehand
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
limit) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            Int
i0 <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Int
is Int
idx
            let run :: Int -> ST s ()
run Int
i = do
                  if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s
                    then do
                      -- within the block
                      MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Bit
block (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) (Bit -> ST s ()) -> Bit -> ST s ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
                      Int -> ST s ()
run (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p
                    else do
                      -- went out of the block
                      MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
is Int
idx Int
i
            Int -> ST s ()
run Int
i0

        Vector Bit
block' <- Int -> Vector Bit -> Vector Bit
forall a. Unbox a => Int -> Vector a -> Vector a
VU.take (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
s (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)) (Vector Bit -> Vector Bit)
-> ST s (Vector Bit) -> ST s (Vector Bit)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState (ST s)) Bit
block
        Vector Bit -> (Int -> Bit -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Bit
block' ((Int -> Bit -> ST s ()) -> ST s ())
-> (Int -> Bit -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i (Bit Bool
b) -> do
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
b (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            Int
at <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
nPrimes Int
0
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
at Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxPrimeCount) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
              MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
nPrimes Int
0 (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
at Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
              MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
result Int
at (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

      Int
len <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
nPrimes Int
0
      MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MVector s Int -> ST s (MVector s Int))
-> MVector s Int -> ST s (MVector s Int)
forall a b. (a -> b) -> a -> b
$ Int -> MVector s Int -> MVector s Int
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VGM.take Int
len MVector s Int
result

-- | \(O(w \log^3 n)\) Miller–Rabin primality test, where \(w = 3\) for \(x \lt 2^{32}\) and
-- \(w = 7\) for \(x \ge 3^{32}\).
--
-- @since 1.2.6.0
{-# INLINEABLE isPrime #-}
isPrime :: Int -> Bool
isPrime :: Int -> Bool
isPrime Int
x
  | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = Bool
False
  -- Up to 11^2:
  | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 Bool -> Bool -> Bool
|| Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3 Bool -> Bool -> Bool
|| Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
5 Bool -> Bool -> Bool
|| Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
7 = Bool
True
  | Int -> Bool
forall a. Integral a => a -> Bool
even Int
x Bool -> Bool -> Bool
|| Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
3 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
5 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
7 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Bool
False
  | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
121 = Bool
True
isPrime Int
x
  -- http://miller-rabin.appspot.com/
  -- \| x < bit 32 = all test [2, 7, 61]
  | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Int
forall a. Bits a => Int -> a
bit Int
32 = Word64 -> Bool
test Word64
2 Bool -> Bool -> Bool
&& Word64 -> Bool
test Word64
7 Bool -> Bool -> Bool
&& Word64 -> Bool
test Word64
61
  -- \| otherwise = all test [2, 325, 9375, 28178, 450775, 9780504, 1795265022]
  | Bool
otherwise = Word64 -> Bool
test Word64
2 Bool -> Bool -> Bool
&& Word64 -> Bool
test Word64
325 Bool -> Bool -> Bool
&& Word64 -> Bool
test Word64
9375 Bool -> Bool -> Bool
&& Word64 -> Bool
test Word64
28178 Bool -> Bool -> Bool
&& Word64 -> Bool
test Word64
450775 Bool -> Bool -> Bool
&& Word64 -> Bool
test Word64
9780504 Bool -> Bool -> Bool
&& Word64 -> Bool
test Word64
1795265022
  where
    !Word64
x64 :: Word64 = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x
    !Word64
d :: Word64 = (Word64
x64 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
.>>. Word64 -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros (Word64
x64 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1)
    !mont :: Montgomery64
mont = Word64 -> Montgomery64
M64.fromVal Word64
x64
    !one :: Word64
one = Montgomery64 -> Word64 -> Word64
M64.encode Montgomery64
mont Word64
1
    !minusOne :: Word64
minusOne = Montgomery64 -> Word64 -> Word64
M64.encode Montgomery64
mont (Word64
x64 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1)
    test :: Word64 -> Bool
test Word64
a = Word64 -> Word64 -> Bool
inner (HasCallStack => Montgomery64 -> Word64 -> Int -> Word64
Montgomery64 -> Word64 -> Int -> Word64
M64.powMod Montgomery64
mont (Montgomery64 -> Word64 -> Word64
M64.encode Montgomery64
mont Word64
a) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
d)) Word64
d
      where
        inner :: Word64 -> Word64 -> Bool
        inner :: Word64 -> Word64 -> Bool
inner Word64
y Word64
t
          | Bool -> Bool
not (Word64 -> Word64 -> Word64 -> Bool
M64.eq Word64
x64 Word64
y Word64
one) Bool -> Bool -> Bool
&& Bool -> Bool
not (Word64 -> Word64 -> Word64 -> Bool
M64.eq Word64
x64 Word64
y Word64
minusOne) Bool -> Bool -> Bool
&& Word64
t Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
x64 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1 = Word64 -> Word64 -> Bool
inner (Montgomery64 -> Word64 -> Word64 -> Word64
M64.mulMod Montgomery64
mont Word64
y Word64
y) (Word64
t Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
.<<. Int
1)
          | Bool -> Bool
not (Word64 -> Word64 -> Word64 -> Bool
M64.eq Word64
x64 Word64
y Word64
minusOne) Bool -> Bool -> Bool
&& Word64 -> Bool
forall a. Integral a => a -> Bool
even Word64
t = Bool
False
          | Bool
otherwise = Bool
True

-- | Pollard's Rho algorithm.
{-# INLINEABLE rho #-}
rho :: (HasCallStack) => Word64 -> Int -> Int -> Int
rho :: HasCallStack => Word64 -> Int -> Int -> Int
rho Word64
modVal Int
n Int
c
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = String -> Int
forall a. HasCallStack => String -> a
error (String -> Int) -> String -> Int
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.Math.rho: given value less than or equal to `1`: `" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
forall a. Show a => a -> String
show String
"`"
  | Bool
otherwise = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Int) -> Word64 -> Int
forall a b. (a -> b) -> a -> b
$! Int -> Word64 -> Word64 -> Word64 -> Word64 -> Word64 -> Word64
inner Int
1 (Montgomery64 -> Word64 -> Word64
M64.encode Montgomery64
mont Word64
1) (Montgomery64 -> Word64 -> Word64
M64.encode Montgomery64
mont Word64
2) (Montgomery64 -> Word64 -> Word64
M64.encode Montgomery64
mont Word64
1) (Montgomery64 -> Word64 -> Word64
M64.encode Montgomery64
mont Word64
1) Word64
1
  where
    -- what a mess!!
    !mont :: Montgomery64
mont = Word64 -> Montgomery64
M64.fromVal Word64
modVal
    !Word64
n64 :: Word64 = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
    !cc :: Word64
cc = Montgomery64 -> Word64 -> Word64
M64.encode Montgomery64
mont (Word64 -> Word64) -> Word64 -> Word64
forall a b. (a -> b) -> a -> b
$ Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c
    f :: Word64 -> Word64
f !Word64
x = Word64 -> Word64 -> Word64 -> Word64
M64.addMod Word64
modVal (Montgomery64 -> Word64 -> Word64 -> Word64
M64.mulMod Montgomery64
mont Word64
x Word64
x) Word64
cc
    fn :: t -> Word64 -> Word64
fn t
0 !Word64
x = Word64
x
    fn t
n_ !Word64
x = t -> Word64 -> Word64
fn (t
n_ t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (Word64 -> Word64) -> Word64 -> Word64
forall a b. (a -> b) -> a -> b
$! Word64 -> Word64
f Word64
x
    !Int
m2 :: Int = Int -> Int
forall a. Bits a => Int -> a
bit (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase (Double
2.0 :: Double) (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
5
    inner :: Int -> Word64 -> Word64 -> Word64 -> Word64 -> Word64 -> Word64
inner Int
r Word64
_lastY0 Word64
y0 Word64
z0 Word64
q0 Word64
g0
      | Word64
g0 Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
1 =
          let !y :: Word64
y = Int -> Word64 -> Word64
forall {t}. (Eq t, Num t) => t -> Word64 -> Word64
fn Int
r Word64
y0
              (!Word64
y', !Word64
z', !Word64
q', !Word64
g') = Int
-> Word64
-> Word64
-> Word64
-> Word64
-> (Word64, Word64, Word64, Word64)
inner2 Int
0 Word64
y Word64
z0 Word64
q0 Word64
g0
           in Int -> Word64 -> Word64 -> Word64 -> Word64 -> Word64 -> Word64
inner (Int
r Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.<<. Int
1) Word64
y0 Word64
y' Word64
z' Word64
q' Word64
g'
      -- FIXME: It can sometimes slow, depending on the seed value
      -- \| g0 == n64 = inner3 z0
      | Bool
otherwise = Word64
g0
      where
        inner2 :: Int
-> Word64
-> Word64
-> Word64
-> Word64
-> (Word64, Word64, Word64, Word64)
inner2 !Int
k !Word64
y !Word64
z !Word64
q !Word64
g
          | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
r Bool -> Bool -> Bool
|| Word64
g Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
1 = (Word64
y, Word64
z, Word64
q, Word64
g)
          | Bool
otherwise =
              let (!Word64
y', !Word64
q') = Int -> Word64 -> Word64 -> (Word64, Word64)
forall {t}.
(Eq t, Num t) =>
t -> Word64 -> Word64 -> (Word64, Word64)
fn2 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m2 (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k)) Word64
y Word64
q
                  !g' :: Word64
g' = Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
gcd (Montgomery64 -> Word64 -> Word64
M64.decode Montgomery64
mont Word64
q) Word64
n64
               in Int
-> Word64
-> Word64
-> Word64
-> Word64
-> (Word64, Word64, Word64, Word64)
inner2 (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
m2) Word64
y' Word64
y Word64
q' Word64
g'
          where
            fn2 :: t -> Word64 -> Word64 -> (Word64, Word64)
fn2 t
0 !Word64
y_ !Word64
q_ = (Word64
y_, Word64
q_)
            fn2 t
n_ !Word64
y_ !Word64
q_ =
              let !y' :: Word64
y' = Word64 -> Word64
f Word64
y_
                  !q' :: Word64
q' = Montgomery64 -> Word64 -> Word64 -> Word64
M64.mulMod Montgomery64
mont Word64
q_ (Word64 -> Word64 -> Word64 -> Word64
M64.subMod Word64
modVal Word64
y0 Word64
y')
               in t -> Word64 -> Word64 -> (Word64, Word64)
fn2 (t
n_ t -> t -> t
forall a. Num a => a -> a -> a
- t
1) Word64
y' Word64
q'

-- FIXME: it can sometimes slow, depending on the seed value
-- inner3 !z
--   | g == 1 = inner3 z'
--   | otherwise = g
--   where
--     !z' = f z
--     !g = gcd (M64.decode mont (M64.subMod modVal lastY0 z')) n64

-- | Tries to find a prime factor for the given value, running Pollard's Rho algorithm.
--
-- ==== Constrants
-- - \(x \gt 1\)
{-# INLINEABLE findPrimeFactor #-}
findPrimeFactor :: (HasCallStack) => StdGen -> Int -> (Maybe Int, StdGen)
findPrimeFactor :: HasCallStack => StdGen -> Int -> (Maybe Int, StdGen)
findPrimeFactor StdGen
gen0 Int
n
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = String -> (Maybe Int, StdGen)
forall a. HasCallStack => String -> a
error (String -> (Maybe Int, StdGen)) -> String -> (Maybe Int, StdGen)
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.Math.findPrimeFactor: given value less than or equal to `1`: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
forall a. Show a => a -> String
show String
"`"
  | Int -> Bool
isPrime Int
n = (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
n, StdGen
gen0)
  | Bool
otherwise = Int -> StdGen -> (Maybe Int, StdGen)
tryN Int
200 StdGen
gen0
  where
    tryN :: Int -> StdGen -> (Maybe Int, StdGen)
    tryN :: Int -> StdGen -> (Maybe Int, StdGen)
tryN Int
0 StdGen
gen = (Maybe Int
forall a. Maybe a
Nothing, StdGen
gen)
    tryN Int
i StdGen
gen
      | Int -> Bool
isPrime Int
m = (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
m, StdGen
gen')
      | Bool
otherwise = Int -> StdGen -> (Maybe Int, StdGen)
tryN (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) StdGen
gen'
      where
        (!Int
rnd, !StdGen
gen') = (Int, Int) -> StdGen -> (Int, StdGen)
forall a g. (UniformRange a, RandomGen g) => (a, a) -> g -> (a, g)
uniformR (Int
0, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) StdGen
gen
        !m :: Int
m = HasCallStack => Word64 -> Int -> Int -> Int
Word64 -> Int -> Int -> Int
rho (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Int
n Int
rnd

-- | Returns prime factors in run-length encoding \((p_i, n_i)\), sorted by \(p_i\).
--
-- ==== Constraints
-- - \(x \ge 1\)
--
-- @since 1.2.6.0
{-# INLINE primeFactors #-}
primeFactors :: (HasCallStack) => Int -> VU.Vector (Int, Int)
primeFactors :: HasCallStack => Int -> Vector (Int, Int)
primeFactors = (forall s. MVector s (Int, Int) -> ST s ())
-> Vector (Int, Int) -> Vector (Int, Int)
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify MVector s (Int, Int) -> ST s ()
MVector (PrimState (ST s)) (Int, Int) -> ST s ()
forall s. MVector s (Int, Int) -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> m ()
VAI.sort (Vector (Int, Int) -> Vector (Int, Int))
-> (Int -> Vector (Int, Int)) -> Int -> Vector (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => Int -> Vector (Int, Int)
Int -> Vector (Int, Int)
primeFactorsUnsorted

-- | Returns prime factors in run-length encoding \((p_i, n_i)\) in arbitrary order.
--
-- It internally uses probabilistic method (Pollard's rho algorithm) and it can actually result in
-- runtime error, however, the probability is very low and the API does not return `Maybe`.
--
-- @since 1.2.6.0
{-# INLINEABLE primeFactorsUnsorted #-}
primeFactorsUnsorted :: (HasCallStack) => Int -> VU.Vector (Int, Int)
primeFactorsUnsorted :: HasCallStack => Int -> Vector (Int, Int)
primeFactorsUnsorted Int
n
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = String -> Vector (Int, Int)
forall a. HasCallStack => String -> a
error (String -> Vector (Int, Int)) -> String -> Vector (Int, Int)
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.Math.primeFactorsUnsorted: given non-positive value `" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"`"
  | Bool
otherwise = (forall s. ST s (MVector s (Int, Int))) -> Vector (Int, Int)
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s (Int, Int))) -> Vector (Int, Int))
-> (forall s. ST s (MVector s (Int, Int))) -> Vector (Int, Int)
forall a b. (a -> b) -> a -> b
$ do
      MVector (PrimState (ST s)) (Int, Int)
buf <- Int -> ST s (MVector (PrimState (ST s)) (Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase (Double
2 :: Double) (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)))

      -- for small prime factors, try them all:
      let runDiv :: Int -> Int -> [Int] -> ST s (Int, Int)
runDiv Int
cur Int
iWrite [] = (Int, Int) -> ST s (Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
cur, Int
iWrite)
          runDiv Int
cur Int
iWrite (Int
d : [Int]
rest)
            | Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
cur = (Int, Int) -> ST s (Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
cur, Int
iWrite)
            | Bool
otherwise = case Int -> Int -> Int -> Maybe (Int, Int)
tryDiv Int
0 Int
cur Int
d of
                Just (!Int
cur', !Int
nd) -> do
                  MVector (PrimState (ST s)) (Int, Int)
-> Int -> (Int, Int) -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) (Int, Int)
buf Int
iWrite (Int
d, Int
nd)
                  Int -> Int -> [Int] -> ST s (Int, Int)
runDiv Int
cur' (Int
iWrite Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int]
rest
                Maybe (Int, Int)
Nothing -> Int -> Int -> [Int] -> ST s (Int, Int)
runDiv Int
cur Int
iWrite [Int]
rest

      (!Int
n', !Int
iWrite0) <- Int -> Int -> [Int] -> ST s (Int, Int)
runDiv Int
n Int
0 (Int
2 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int
3, Int
5 .. Int
97])

      -- for bigger prime numbers, use Polland's rho algorithm:
      let runRho :: StdGen -> Int -> Int -> m Int
runRho !StdGen
gen !Int
cur !Int
iWrite
            | Int
cur Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 = case HasCallStack => StdGen -> Int -> (Maybe Int, StdGen)
StdGen -> Int -> (Maybe Int, StdGen)
findPrimeFactor StdGen
gen Int
cur of
                (Just Int
p, !StdGen
gen') -> do
                  let (!Int
cur', !Int
np) = Maybe (Int, Int) -> (Int, Int)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Int, Int) -> (Int, Int)) -> Maybe (Int, Int) -> (Int, Int)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Maybe (Int, Int)
tryDiv Int
0 Int
cur Int
p
                  MVector (PrimState m) (Int, Int) -> Int -> (Int, Int) -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) (Int, Int)
MVector (PrimState (ST s)) (Int, Int)
buf Int
iWrite (Int
p, Int
np)
                  StdGen -> Int -> Int -> m Int
runRho StdGen
gen' Int
cur' (Int
iWrite Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                (Maybe Int
Nothing, !StdGen
_gen') -> do
                  -- we could return `Nothing` instead
                  String -> m Int
forall a. HasCallStack => String -> a
error (String -> m Int) -> String -> m Int
forall a b. (a -> b) -> a -> b
$ String
"unable to find prime factor for " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
cur
            | Bool
otherwise = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
iWrite

      -- NOTE: The seed value ifs fixed here. We could decide it at runtime for possibly faster
      -- submissions on TLE redjuge, however, 're rather preferring deterministic result:
      Int
len <- StdGen -> Int -> Int -> ST s Int
forall {m :: * -> *}.
(PrimState m ~ s, PrimMonad m) =>
StdGen -> Int -> Int -> m Int
runRho (Int -> StdGen
mkStdGen Int
123456789) Int
n' Int
iWrite0
      MVector s (Int, Int) -> ST s (MVector s (Int, Int))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MVector s (Int, Int) -> ST s (MVector s (Int, Int)))
-> MVector s (Int, Int) -> ST s (MVector s (Int, Int))
forall a b. (a -> b) -> a -> b
$ Int -> MVector s (Int, Int) -> MVector s (Int, Int)
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.take Int
len MVector s (Int, Int)
MVector (PrimState (ST s)) (Int, Int)
buf
  where
    tryDiv :: Int -> Int -> Int -> Maybe (Int, Int)
    tryDiv :: Int -> Int -> Int -> Maybe (Int, Int)
tryDiv !Int
nDiv Int
x Int
d
      | Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> Int -> Int -> Maybe (Int, Int)
tryDiv (Int
nDiv Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
q Int
d
      | Int
nDiv Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = (Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
x, Int
nDiv)
      | Bool
otherwise = Maybe (Int, Int)
forall a. Maybe a
Nothing
      where
        (!Int
q, !Int
r) = Int
x Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
d

-- | Enumerates divisors of the input value.
--
-- ==== Constraints
-- - \(x \ge 1\)
--
-- @since 1.2.6.0
{-# INLINE divisors #-}
divisors :: Int -> VU.Vector Int
-- TODO: use intro sort?
divisors :: Int -> Vector Int
divisors = (forall s. MVector s Int -> ST s ()) -> Vector Int -> Vector Int
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify MVector s Int -> ST s ()
MVector (PrimState (ST s)) Int -> ST s ()
forall s. MVector s Int -> ST s ()
forall e (m :: * -> *) (v :: * -> * -> *).
(PrimMonad m, MVector v e, Radix e) =>
v (PrimState m) e -> m ()
VAR.sort (Vector Int -> Vector Int)
-> (Int -> Vector Int) -> Int -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Vector Int
divisorsUnsorted

-- | Enumerates divisors of the input value.
--
-- ==== Constraints
-- - \(x \ge 1\)
--
-- @since 1.2.6.0
{-# INLINEABLE divisorsUnsorted #-}
divisorsUnsorted :: Int -> VU.Vector Int
divisorsUnsorted :: Int -> Vector Int
divisorsUnsorted Int
x = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
  MVector s Int
vec <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
nDivisors
  MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
vec Int
0 Int
1
  (Int -> (Int, Int) -> ST s Int)
-> Int -> Vector (Int, Int) -> ST s ()
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m ()
VU.foldM'_
    ( \Int
lenSofar (!Int
p, !Int
np) -> do
        ((Int, Int) -> Int
forall a b. (a, b) -> a
fst <$>)
          (ST s (Int, Int) -> ST s Int) -> ST s (Int, Int) -> ST s Int
forall a b. (a -> b) -> a -> b
$ ((Int, Int) -> Int -> ST s (Int, Int))
-> (Int, Int) -> Vector Int -> ST s (Int, Int)
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM'
            ( \(!Int
offset, !Int
pp) Int
_ -> do
                let !pp' :: Int
pp' = Int
pp Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p
                -- multiply to all the values sofar:
                MVector (PrimState (ST s)) Int
-> (Int -> Int -> ST s Int) -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a b.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (Int -> a -> m b) -> m ()
VGM.iforM_ (Int -> MVector s Int -> MVector s Int
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VGM.take Int
lenSofar MVector s Int
vec) ((Int -> Int -> ST s Int) -> ST s ())
-> (Int -> Int -> ST s Int) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i Int
vx -> do
                  MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
vec (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int
vx Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
pp'
                  Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
pp'
                (Int, Int) -> ST s (Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lenSofar, Int
pp')
            )
            (Int
lenSofar, Int
1 :: Int)
          (Vector Int -> ST s (Int, Int)) -> Vector Int -> ST s (Int, Int)
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
np (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    )
    (Int
1 :: Int)
    Vector (Int, Int)
pns
  MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
vec
  where
    pns :: Vector (Int, Int)
pns = HasCallStack => Int -> Vector (Int, Int)
Int -> Vector (Int, Int)
primeFactors Int
x
    (!Vector Int
_, !Vector Int
ns) = Vector (Int, Int) -> (Vector Int, Vector Int)
forall a b.
(Unbox a, Unbox b) =>
Vector (a, b) -> (Vector a, Vector b)
VU.unzip Vector (Int, Int)
pns
    nDivisors :: Int
nDivisors = (Int -> Int -> Int) -> Int -> Vector Int -> Int
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' (\ !Int
acc Int
n -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Int
1 :: Int) Vector Int
ns

-- | Calculates \(x^n\) with custom multiplication operator using the binary exponentiation
-- technique.
--
-- The internal implementation is taken from @Data.Semigroup.stimes@, but `power` uses strict
-- evaluation and is often much faster.
--
-- ==== Complexity
-- - \(O(\log n)\)
--
-- ==== Constraints
-- - \(n \gt 0\)
--
-- @since 1.0.0.0
{-# INLINE power #-}
power :: (a -> a -> a) -> Int -> a -> a
power :: forall a. (a -> a -> a) -> Int -> a -> a
power a -> a -> a
op Int
n0 a
x1
  | Int
n0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = String -> a
forall a. String -> a
errorWithoutStackTrace String
"AtCoder.Extra.Math.power: positive multiplier expected"
  | Bool
otherwise = a -> Int -> a
forall {a}. (Integral a, Bits a) => a -> a -> a
f a
x1 Int
n0
  where
    f :: a -> a -> a
f !a
x !a
n
      | a -> Bool
forall a. Integral a => a -> Bool
even a
n = a -> a -> a
f (a
x a -> a -> a
`op` a
x) (a
n a -> Int -> a
forall a. Bits a => a -> Int -> a
.>>. Int
1)
      | a
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 = a
x
      | Bool
otherwise = a -> a -> a -> a
forall {a}. (Integral a, Bits a) => a -> a -> a -> a
g (a
x a -> a -> a
`op` a
x) (a
n a -> Int -> a
forall a. Bits a => a -> Int -> a
.>>. Int
1) a
x
    g :: a -> a -> a -> a
g !a
x !a
n !a
z
      | a -> Bool
forall a. Integral a => a -> Bool
even a
n = a -> a -> a -> a
g (a
x a -> a -> a
`op` a
x) (a
n a -> Int -> a
forall a. Bits a => a -> Int -> a
.>>. Int
1) a
z
      | a
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 = a
x a -> a -> a
`op` a
z
      | Bool
otherwise = a -> a -> a -> a
g (a
x a -> a -> a
`op` a
x) (a
n a -> Int -> a
forall a. Bits a => a -> Int -> a
.>>. Int
1) (a
x a -> a -> a
`op` a
z)

-- | Strict variant of @Data.Semigroup.stimes@.
--
-- ==== Complexity
-- - \(O(\log n)\)
--
-- ==== Constraints
-- - \(n \gt 0\)
--
-- @since 1.0.0.0
{-# INLINE stimes' #-}
stimes' :: (Semigroup a) => Int -> a -> a
stimes' :: forall a. Semigroup a => Int -> a -> a
stimes' = (a -> a -> a) -> Int -> a -> a
forall a. (a -> a -> a) -> Int -> a -> a
power a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>)

-- | Strict variant of @Data.Monoid.mtimes@.
--
-- ==== Complexity
-- - \(O(\log n)\)
--
-- ==== Constraints
-- - \(n \ge 0\)
--
-- @since 1.0.0.0
{-# INLINE mtimes' #-}
mtimes' :: (Monoid a) => Int -> a -> a
mtimes' :: forall a. Monoid a => Int -> a -> a
mtimes' Int
n a
x = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
n Int
0 of
  Ordering
LT -> String -> a
forall a. String -> a
errorWithoutStackTrace String
"AtCoder.Extra.Math.mtimes': non-negative multiplier expected"
  Ordering
EQ -> a
forall a. Monoid a => a
mempty
  Ordering
GT -> (a -> a -> a) -> Int -> a -> a
forall a. (a -> a -> a) -> Int -> a -> a
power a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>) Int
n a
x