{-# 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 :: Proxy v -> Word -> t
totientSum Proxy v
_ Word
0 = t
0
totientSum Proxy v
proxy Word
n = Proxy v -> (Word -> t) -> Word -> t
forall t (v :: * -> *).
(Num t, Vector v t) =>
Proxy v -> (Word -> t) -> Word -> t
generalInversion Proxy v
proxy (t -> t
forall a. Integral a => a -> a
triangle (t -> t) -> (Word -> t) -> Word -> t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral) Word
n
where
triangle :: a -> a
triangle a
k = (a
k a -> a -> a
forall a. Num a => a -> a -> a
* (a
k a -> a -> a
forall a. Num a => a -> a -> a
+ a
1)) a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
2
generalInversion
:: (Num t, G.Vector v t)
=> Proxy v
-> (Word -> t)
-> Word
-> t
generalInversion :: Proxy v -> (Word -> t) -> Word -> t
generalInversion Proxy v
proxy Word -> t
fun Word
n = case Word
n of
Word
0 ->[Char] -> t
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 t -> t -> t
forall a. Num a => a -> a -> a
- Word -> t
fun Word
1
Word
3 -> Word -> t
fun Word
3 t -> t -> t
forall a. Num a => a -> a -> a
- t
2t -> t -> t
forall a. Num a => a -> a -> a
*Word -> t
fun Word
1
Word
_ -> (forall s. ST s t) -> t
forall a. (forall s. ST s a) -> a
runST (Proxy v -> (Int -> t) -> Int -> ST s t
forall s t (v :: * -> *).
(Num t, Vector v t) =>
Proxy v -> (Int -> t) -> Int -> ST s t
fastInvertST Proxy v
proxy (Word -> t
fun (Word -> t) -> (Int -> Word) -> Int -> t
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 :: Proxy v -> (Int -> t) -> Int -> ST s t
fastInvertST Proxy v
_ Int -> t
fun Int
n = do
let !k0 :: Int
k0 = Int -> Int
forall a. Integral a => a -> a
integerSquareRoot (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2)
!mk0 :: Int
mk0 = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
k0Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
kmax :: a -> a -> a
kmax a
a a
m = (a
a a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
m a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
2
Mutable v s t
small <- Int -> ST s (Mutable v (PrimState (ST s)) t)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew (Int
mk0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) :: ST s (G.Mutable v s t)
Mutable v (PrimState (ST s)) t -> Int -> t -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
Mutable v (PrimState (ST s)) t
small Int
0 t
0
Mutable v (PrimState (ST s)) t -> Int -> t -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
Mutable v (PrimState (ST s)) t
small Int
1 (t -> ST s ()) -> t -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int -> t
fun Int
1
Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
mk0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
Mutable v (PrimState (ST s)) t -> Int -> t -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
Mutable v (PrimState (ST s)) t
small Int
2 (t -> ST s ()) -> t -> ST s ()
forall a b. (a -> b) -> a -> b
$! (Int -> t
fun Int
2 t -> t -> t
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i = (Int, Int) -> ST s (Int, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
switch,Int
change)
| Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
change = Int -> Int -> Int -> ST s (Int, Int)
calcit (Int
switchInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
change Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
switchInt -> Int -> Int
forall 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 Int -> Int -> Bool
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 <- Mutable v (PrimState (ST s)) t -> Int -> ST s t
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
Mutable v (PrimState (ST s)) t
small Int
m
let nxtk :: Int
nxtk = Int -> Int -> Int
forall a. Integral a => a -> a -> a
kmax Int
i (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
t -> Int -> Int -> ST s (Int, Int)
mloop (t
acc t -> t -> t
forall a. Num a => a -> a -> a
- Int -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
nxtk)t -> t -> t
forall a. Num a => a -> a -> a
*t
val) Int
nxtk (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
kloop :: t -> Int -> ST s (Int, Int)
kloop !t
acc Int
k
| Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = do
Mutable v (PrimState (ST s)) t -> Int -> t -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
Mutable v (PrimState (ST s)) t
small Int
i (t -> ST s ()) -> t -> ST s ()
forall a b. (a -> b) -> a -> b
$! t
acc
Int -> Int -> Int -> ST s (Int, Int)
calcit Int
switch Int
change (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
| Bool
otherwise = do
t
val <- Mutable v (PrimState (ST s)) t -> Int -> ST s t
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
Mutable v (PrimState (ST s)) t
small (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))
t -> Int -> ST s (Int, Int)
kloop (t
acct -> t -> t
forall a. Num a => a -> a -> a
-t
val) (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
t -> Int -> Int -> ST s (Int, Int)
mloop (Int -> t
fun Int
i t -> t -> t
forall a. Num a => a -> a -> a
- Int -> t
fun (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2)) ((Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int -> Int -> Int
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 <- Int -> ST s (Mutable v (PrimState (ST s)) t)
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Mutable v s t -> ST s (Mutable v s t)
forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v s t
large
| (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
change Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n = Int -> Int -> Int -> ST s (Mutable v s t)
calcbig (Int
switchInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
change Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
switchInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
6) Int
j
| Bool
otherwise = do
let i :: Int
i = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
jInt -> Int -> Int
forall 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 Int -> Int -> Bool
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 <- Mutable v (PrimState (ST s)) t -> Int -> ST s t
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
Mutable v (PrimState (ST s)) t
small Int
m
let nxtk :: Int
nxtk = Int -> Int -> Int
forall a. Integral a => a -> a -> a
kmax Int
i (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
t -> Int -> Int -> ST s (Mutable v s t)
mloop (t
acc t -> t -> t
forall a. Num a => a -> a -> a
- Int -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
nxtk)t -> t -> t
forall a. Num a => a -> a -> a
*t
val) Int
nxtk (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
kloop :: t -> Int -> ST s (Mutable v s t)
kloop !t
acc Int
k
| Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = do
Mutable v (PrimState (ST s)) t -> Int -> t -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
Mutable v (PrimState (ST s)) t
large (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (t -> ST s ()) -> t -> ST s ()
forall a b. (a -> b) -> a -> b
$! t
acc
Int -> Int -> Int -> ST s (Mutable v s t)
calcbig Int
switch Int
change (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
| Bool
otherwise = do
let m :: Int
m = Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
t
val <- if Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
mk0
then Mutable v (PrimState (ST s)) t -> Int -> ST s t
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
Mutable v (PrimState (ST s)) t
small Int
m
else Mutable v (PrimState (ST s)) t -> Int -> ST s t
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
Mutable v (PrimState (ST s)) t
large (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
*(Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
t -> Int -> ST s (Mutable v s t)
kloop (t
acct -> t -> t
forall a. Num a => a -> a -> a
-t
val) (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
t -> Int -> Int -> ST s (Mutable v s t)
mloop (Int -> t
fun Int
i t -> t -> t
forall a. Num a => a -> a -> a
- Int -> t
fun (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2)) ((Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int -> Int -> Int
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
Mutable v (PrimState (ST s)) t -> Int -> ST s t
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
Mutable v (PrimState (ST s)) t
mvec Int
0