1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
|
-- |
-- Module: Math.NumberTheory.Utils.DirichletSeries
-- Copyright: (c) 2018 Andrew Lelechenko
-- Licence: MIT
-- Maintainer: Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- An abstract representation of a Dirichlet series over a semiring.
--
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
module Math.NumberTheory.Utils.DirichletSeries
( DirichletSeries
, fromDistinctAscList
, lookup
, filter
, partition
, unions
, union
, size
, timesAndCrop
) where
import Prelude hiding (filter, last, rem, quot, snd, lookup)
import Data.Coerce
import Data.Euclidean
import Data.Map (Map)
import qualified Data.Map.Strict as M
import Data.Maybe
import Data.Semiring (Semiring(..))
import Numeric.Natural
-- Sparse Dirichlet series are represented by an ascending list of pairs.
-- For instance, [(a, b), (c, d)] stands for 1 + b/a^s + d/c^s.
-- Note that the representation still may include a term (1, b), so
-- [(1, b), (c, d)] means (1 + b) + d/c^s.
newtype DirichletSeries a b = DirichletSeries { _unDirichletSeries :: Map a b }
deriving (Show)
fromDistinctAscList :: forall a b. [(a, b)] -> DirichletSeries a b
fromDistinctAscList = coerce (M.fromDistinctAscList @a @b)
lookup :: (Ord a, Num a, Semiring b) => a -> DirichletSeries a b -> b
lookup 1 (DirichletSeries m) = M.findWithDefault zero 1 m `plus` one
lookup a (DirichletSeries m) = M.findWithDefault zero a m
filter :: forall a b. (a -> Bool) -> DirichletSeries a b -> DirichletSeries a b
filter predicate = coerce (M.filterWithKey @a @b (\k _ -> predicate k))
partition :: forall a b. (a -> Bool) -> DirichletSeries a b -> (DirichletSeries a b, DirichletSeries a b)
partition predicate = coerce (M.partitionWithKey @a @b (\k _ -> predicate k))
unions :: forall a b. (Ord a, Semiring b) => [DirichletSeries a b] -> DirichletSeries a b
unions = coerce (M.unionsWith plus :: [Map a b] -> Map a b)
union :: forall a b. (Ord a, Semiring b) => DirichletSeries a b -> DirichletSeries a b -> DirichletSeries a b
union = coerce (M.unionWith @a @b plus)
size :: forall a b. DirichletSeries a b -> Int
size = coerce (M.size @a @b)
-- | Let as = sum_i k_i/a_i^s and bs = sum_i l_i/b_i^s be Dirichlet series,
-- and all a_i and b_i are divisors of n. Return Dirichlet series cs,
-- which contains all terms as * bs = sum_i m_i/c_i^s such that c_i divides n.
timesAndCrop
:: (Num a, Euclidean a, Ord a, Semiring b)
=> a
-> DirichletSeries a b
-> DirichletSeries a b
-> DirichletSeries a b
timesAndCrop n (DirichletSeries as) (DirichletSeries bs)
= DirichletSeries
$ M.unionWith plus (M.unionWith plus as bs)
$ M.fromListWith plus
[ (a * b, fa `times` fb)
| (b, fb) <- M.assocs bs
, let nb = n `quot` b
, (a, fa) <- takeWhile ((<= nb) . fst) (M.assocs as)
, isJust (nb `divide` a)
]
{-# SPECIALISE timesAndCrop :: Semiring b => Int -> DirichletSeries Int b -> DirichletSeries Int b -> DirichletSeries Int b #-}
{-# SPECIALISE timesAndCrop :: Semiring b => Word -> DirichletSeries Word b -> DirichletSeries Word b -> DirichletSeries Word b #-}
{-# SPECIALISE timesAndCrop :: Semiring b => Integer -> DirichletSeries Integer b -> DirichletSeries Integer b -> DirichletSeries Integer b #-}
{-# SPECIALISE timesAndCrop :: Semiring b => Natural -> DirichletSeries Natural b -> DirichletSeries Natural b -> DirichletSeries Natural b #-}
|
