{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Math.NumberTheory.MoebiusInversion
( generalInversion
, totientSum
) where
import Control.Monad
import Control.Monad.ST
import Data.Proxy
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as MG
import Math.NumberTheory.Roots
import Math.NumberTheory.Utils.FromIntegral
totientSum
:: (Integral t, G.Vector v t)
=> Proxy v
-> Word
-> t
totientSum :: forall t (v :: * -> *).
(Integral t, Vector v t) =>
Proxy v -> Word -> t
totientSum Proxy v
_ Word
0 = t
0
totientSum Proxy v
proxy Word
n = forall t (v :: * -> *).
(Num t, Vector v t) =>
Proxy v -> (Word -> t) -> Word -> t
generalInversion Proxy v
proxy (forall {a}. Integral a => a -> a
triangle forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral) Word
n
where
triangle :: a -> a
triangle a
k = (a
k forall a. Num a => a -> a -> a
* (a
k forall a. Num a => a -> a -> a
+ a
1)) forall a. Integral a => a -> a -> a
`quot` a
2
generalInversion
:: (Num t, G.Vector v t)
=> Proxy v
-> (Word -> t)
-> Word
-> t
generalInversion :: forall t (v :: * -> *).
(Num t, Vector v t) =>
Proxy v -> (Word -> t) -> Word -> t
generalInversion Proxy v
proxy Word -> t
fun Word
n = case Word
n of
Word
0 ->forall a. HasCallStack => [Char] -> a
error [Char]
"Möbius inversion only defined on positive domain"
Word
1 -> Word -> t
fun Word
1
Word
2 -> Word -> t
fun Word
2 forall a. Num a => a -> a -> a
- Word -> t
fun Word
1
Word
3 -> Word -> t
fun Word
3 forall a. Num a => a -> a -> a
- t
2forall a. Num a => a -> a -> a
*Word -> t
fun Word
1
Word
_ -> forall a. (forall s. ST s a) -> a
runST (forall s t (v :: * -> *).
(Num t, Vector v t) =>
Proxy v -> (Int -> t) -> Int -> ST s t
fastInvertST Proxy v
proxy (Word -> t
fun forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word
intToWord) (Word -> Int
wordToInt Word
n))
fastInvertST
:: forall s t v.
(Num t, G.Vector v t)
=> Proxy v
-> (Int -> t)
-> Int
-> ST s t
fastInvertST :: forall s t (v :: * -> *).
(Num t, Vector v t) =>
Proxy v -> (Int -> t) -> Int -> ST s t
fastInvertST Proxy v
_ Int -> t
fun Int
n = do
let !k0 :: Int
k0 = forall {a}. Integral a => a -> a
integerSquareRoot (Int
n forall a. Integral a => a -> a -> a
`quot` Int
2)
!mk0 :: Int
mk0 = Int
n forall a. Integral a => a -> a -> a
`quot` (Int
2forall a. Num a => a -> a -> a
*Int
k0forall a. Num a => a -> a -> a
+Int
1)
kmax :: a -> a -> a
kmax a
a a
m = (a
a forall a. Integral a => a -> a -> a
`quot` a
m forall a. Num a => a -> a -> a
- a
1) forall a. Integral a => a -> a -> a
`quot` a
2
Mutable v s t
small <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew (Int
mk0 forall a. Num a => a -> a -> a
+ Int
1) :: ST s (G.Mutable v s t)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
small Int
0 t
0
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
small Int
1 forall a b. (a -> b) -> a -> b
$! Int -> t
fun Int
1
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
mk0 forall a. Ord a => a -> a -> Bool
>= Int
2) forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
small Int
2 forall a b. (a -> b) -> a -> b
$! (Int -> t
fun Int
2 forall a. Num a => a -> a -> a
- Int -> t
fun Int
1)
let calcit :: Int -> Int -> Int -> ST s (Int, Int)
calcit :: Int -> Int -> Int -> ST s (Int, Int)
calcit Int
switch Int
change Int
i
| Int
mk0 forall a. Ord a => a -> a -> Bool
< Int
i = forall (m :: * -> *) a. Monad m => a -> m a
return (Int
switch,Int
change)
| Int
i forall a. Eq a => a -> a -> Bool
== Int
change = Int -> Int -> Int -> ST s (Int, Int)
calcit (Int
switchforall a. Num a => a -> a -> a
+Int
1) (Int
change forall a. Num a => a -> a -> a
+ Int
4forall a. Num a => a -> a -> a
*Int
switchforall a. Num a => a -> a -> a
+Int
6) Int
i
| Bool
otherwise = do
let mloop :: t -> Int -> Int -> ST s (Int, Int)
mloop !t
acc Int
k !Int
m
| Int
k forall a. Ord a => a -> a -> Bool
< Int
switch = t -> Int -> ST s (Int, Int)
kloop t
acc Int
k
| Bool
otherwise = do
t
val <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
small Int
m
let nxtk :: Int
nxtk = forall a. Integral a => a -> a -> a
kmax Int
i (Int
mforall a. Num a => a -> a -> a
+Int
1)
t -> Int -> Int -> ST s (Int, Int)
mloop (t
acc forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
kforall a. Num a => a -> a -> a
-Int
nxtk)forall a. Num a => a -> a -> a
*t
val) Int
nxtk (Int
mforall a. Num a => a -> a -> a
+Int
1)
kloop :: t -> Int -> ST s (Int, Int)
kloop !t
acc Int
k
| Int
k forall a. Eq a => a -> a -> Bool
== Int
0 = do
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
small Int
i forall a b. (a -> b) -> a -> b
$! t
acc
Int -> Int -> Int -> ST s (Int, Int)
calcit Int
switch Int
change (Int
iforall a. Num a => a -> a -> a
+Int
1)
| Bool
otherwise = do
t
val <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
small (Int
i forall a. Integral a => a -> a -> a
`quot` (Int
2forall a. Num a => a -> a -> a
*Int
kforall a. Num a => a -> a -> a
+Int
1))
t -> Int -> ST s (Int, Int)
kloop (t
accforall a. Num a => a -> a -> a
-t
val) (Int
kforall a. Num a => a -> a -> a
-Int
1)
t -> Int -> Int -> ST s (Int, Int)
mloop (Int -> t
fun Int
i forall a. Num a => a -> a -> a
- Int -> t
fun (Int
i forall a. Integral a => a -> a -> a
`quot` Int
2)) ((Int
iforall a. Num a => a -> a -> a
-Int
1) forall a. Integral a => a -> a -> a
`quot` Int
2) Int
1
(Int
sw, Int
ch) <- Int -> Int -> Int -> ST s (Int, Int)
calcit Int
1 Int
8 Int
3
Mutable v s t
large <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew Int
k0 :: ST s (G.Mutable v s t)
let calcbig :: Int -> Int -> Int -> ST s (G.Mutable v s t)
calcbig :: Int -> Int -> Int -> ST s (Mutable v s t)
calcbig Int
switch Int
change Int
j
| Int
j forall a. Eq a => a -> a -> Bool
== Int
0 = forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v s t
large
| (Int
2forall a. Num a => a -> a -> a
*Int
jforall a. Num a => a -> a -> a
-Int
1)forall a. Num a => a -> a -> a
*Int
change forall a. Ord a => a -> a -> Bool
<= Int
n = Int -> Int -> Int -> ST s (Mutable v s t)
calcbig (Int
switchforall a. Num a => a -> a -> a
+Int
1) (Int
change forall a. Num a => a -> a -> a
+ Int
4forall a. Num a => a -> a -> a
*Int
switchforall a. Num a => a -> a -> a
+Int
6) Int
j
| Bool
otherwise = do
let i :: Int
i = Int
n forall a. Integral a => a -> a -> a
`quot` (Int
2forall a. Num a => a -> a -> a
*Int
jforall a. Num a => a -> a -> a
-Int
1)
mloop :: t -> Int -> Int -> ST s (Mutable v s t)
mloop !t
acc Int
k Int
m
| Int
k forall a. Ord a => a -> a -> Bool
< Int
switch = t -> Int -> ST s (Mutable v s t)
kloop t
acc Int
k
| Bool
otherwise = do
t
val <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
small Int
m
let nxtk :: Int
nxtk = forall a. Integral a => a -> a -> a
kmax Int
i (Int
mforall a. Num a => a -> a -> a
+Int
1)
t -> Int -> Int -> ST s (Mutable v s t)
mloop (t
acc forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
kforall a. Num a => a -> a -> a
-Int
nxtk)forall a. Num a => a -> a -> a
*t
val) Int
nxtk (Int
mforall a. Num a => a -> a -> a
+Int
1)
kloop :: t -> Int -> ST s (Mutable v s t)
kloop !t
acc Int
k
| Int
k forall a. Eq a => a -> a -> Bool
== Int
0 = do
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
large (Int
jforall a. Num a => a -> a -> a
-Int
1) forall a b. (a -> b) -> a -> b
$! t
acc
Int -> Int -> Int -> ST s (Mutable v s t)
calcbig Int
switch Int
change (Int
jforall a. Num a => a -> a -> a
-Int
1)
| Bool
otherwise = do
let m :: Int
m = Int
i forall a. Integral a => a -> a -> a
`quot` (Int
2forall a. Num a => a -> a -> a
*Int
kforall a. Num a => a -> a -> a
+Int
1)
t
val <- if Int
m forall a. Ord a => a -> a -> Bool
<= Int
mk0
then forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
small Int
m
else forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
large (Int
kforall a. Num a => a -> a -> a
*(Int
2forall a. Num a => a -> a -> a
*Int
jforall a. Num a => a -> a -> a
-Int
1)forall a. Num a => a -> a -> a
+Int
jforall a. Num a => a -> a -> a
-Int
1)
t -> Int -> ST s (Mutable v s t)
kloop (t
accforall a. Num a => a -> a -> a
-t
val) (Int
kforall a. Num a => a -> a -> a
-Int
1)
t -> Int -> Int -> ST s (Mutable v s t)
mloop (Int -> t
fun Int
i forall a. Num a => a -> a -> a
- Int -> t
fun (Int
i forall a. Integral a => a -> a -> a
`quot` Int
2)) ((Int
iforall a. Num a => a -> a -> a
-Int
1) forall a. Integral a => a -> a -> a
`quot` Int
2) Int
1
Mutable v s t
mvec <- Int -> Int -> Int -> ST s (Mutable v s t)
calcbig Int
sw Int
ch Int
k0
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
mvec Int
0