-- SPDX-License-Identifier: MPL-2.0

{- | Reference numerical search and ordering algorithms. Search routines validate
ascending input (O(n)); their subsequent search is logarithmic except ternary
optimization. Pure comparison sorts accept any lawful total 'Ord'.
-}
module HMath.Algorithms (
    lowerBound,
    upperBound,
    binarySearch,
    interpolationSearch,
    exponentialSearch,
    ternarySearch,
    quickselect,
    introselect,
    partialSort,
    mergeSort,
    quicksort,
    heapsort,
    introsort,
    countingSort,
    radixSort,
    bucketSort,
    argsort,
    topK,
    quantileSelect,
    deduplicate,
    numericGroups,
    runLengthEncode,
    runLengthDecode,
    prefixSums,
    suffixSums,
    differenceArray,
    applyDifferences,
    coordinateCompression,
    intervalSweep,
    divideConquerDP,
    tabulateDP,
    memoizedDP,
    binaryLift,
    integerMatrixPower,
    polynomialHash,
    rollingHashes,
) where

import Data.Array (elems, listArray, (!), (//))
import qualified Data.Foldable as Foldable (foldl')
import Data.List (group, sort, sortOn)
import qualified Data.Map.Strict as M
import HMath.Error (MathError (..))

ordered :: (Ord a) => [a] -> Either MathError ()
ordered :: forall a. Ord a => [a] -> Either MathError ()
ordered [a]
xs
    | [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((a -> a -> Bool) -> [a] -> [a] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=) [a]
xs (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
1 [a]
xs)) = () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    | Bool
otherwise = MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"search input must be ascending")

bound :: (Ord a) => (a -> a -> Bool) -> a -> [a] -> Int
bound :: forall a. Ord a => (a -> a -> Bool) -> a -> [a] -> Int
bound a -> a -> Bool
cmp a
x [a]
xs = Int -> Int -> Int
go Int
0 ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)
  where
    a :: Array Int a
a = (Int, Int) -> [a] -> Array Int a
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0, [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
xs
    go :: Int -> Int -> Int
go Int
lo Int
hi
        | Int
lo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
hi = Int
lo
        | a -> a -> Bool
cmp (Array Int a
a Array Int a -> Int -> a
forall i e. Ix i => Array i e -> i -> e
! Int
mid) a
x = Int -> Int -> Int
go (Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
hi
        | Bool
otherwise = Int -> Int -> Int
go Int
lo Int
mid
      where
        mid :: Int
mid = Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
hi Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lo) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
lowerBound, upperBound :: (Ord a) => a -> [a] -> Either MathError Int

{- | Given a target and ascending list, return the first zero-based insertion index whose
value is >= target, or length at the end. Reject unsorted input. Validation is O(n).
Example: lowerBound 2 [1,2,2,4] == Right 1.
-}
lowerBound :: forall a. Ord a => a -> [a] -> Either MathError Int
lowerBound a
x [a]
xs = [a] -> Either MathError ()
forall a. Ord a => [a] -> Either MathError ()
ordered [a]
xs Either MathError () -> Either MathError Int -> Either MathError Int
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Either MathError Int
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a -> a -> Bool) -> a -> [a] -> Int
forall a. Ord a => (a -> a -> Bool) -> a -> [a] -> Int
bound a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<) a
x [a]
xs)

{- | Given a target and ascending list, return the first insertion index whose value is >
target. Reject unsorted input. Example: upperBound 2 [1,2,2,4] == Right 3.
-}
upperBound :: forall a. Ord a => a -> [a] -> Either MathError Int
upperBound a
x [a]
xs = [a] -> Either MathError ()
forall a. Ord a => [a] -> Either MathError ()
ordered [a]
xs Either MathError () -> Either MathError Int -> Either MathError Int
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Either MathError Int
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a -> a -> Bool) -> a -> [a] -> Int
forall a. Ord a => (a -> a -> Bool) -> a -> [a] -> Int
bound a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=) a
x [a]
xs)

{- | Given target and ascending list, return Just the first matching zero-based index or
Nothing if absent. Reject unsorted input. Example: binarySearch 2 [1,2,2] == Right
(Just 1).
-}
binarySearch :: (Ord a) => a -> [a] -> Either MathError (Maybe Int)
binarySearch :: forall a. Ord a => a -> [a] -> Either MathError (Maybe Int)
binarySearch a
x [a]
xs = do i <- a -> [a] -> Either MathError Int
forall a. Ord a => a -> [a] -> Either MathError Int
lowerBound a
x [a]
xs; pure (if i < length xs && xs !! i == x then Just i else Nothing)

{- | Search ascending Integer values using value-based interpolation, returning the first
matching index or Nothing. Reject unsorted input. Example: interpolationSearch 20
[10,20,30] == Right (Just 1).
-}
interpolationSearch :: Integer -> [Integer] -> Either MathError (Maybe Int)
interpolationSearch :: Integer -> [Integer] -> Either MathError (Maybe Int)
interpolationSearch Integer
x [Integer]
xs = do [Integer] -> Either MathError ()
forall a. Ord a => [a] -> Either MathError ()
ordered [Integer]
xs; Maybe Int -> Either MathError (Maybe Int)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Int -> Maybe Int
go Int
0 ([Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
  where
    a :: Array Int Integer
a = (Int, Int) -> [Integer] -> Array Int Integer
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0, [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Integer]
xs
    go :: Int -> Int -> Maybe Int
go Int
lo Int
hi
        | Int
lo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
hi Bool -> Bool -> Bool
|| Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Array Int Integer
a Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
lo Bool -> Bool -> Bool
|| Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Array Int Integer
a Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
hi = Maybe Int
forall a. Maybe a
Nothing
        | Array Int Integer
a Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
lo Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Array Int Integer
a Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
hi = if Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Array Int Integer
a Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
lo then Int -> Maybe Int
forall a. a -> Maybe a
Just Int
lo else Maybe Int
forall a. Maybe a
Nothing
        | Array Int Integer
a Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
p Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
x = Int -> Maybe Int
forall a. a -> Maybe a
Just ((Integer -> Integer -> Bool) -> Integer -> [Integer] -> Int
forall a. Ord a => (a -> a -> Bool) -> a -> [a] -> Int
bound Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
(<) Integer
x [Integer]
xs)
        | Array Int Integer
a Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
p Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
x = Int -> Int -> Maybe Int
go (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
hi
        | Bool
otherwise = Int -> Int -> Maybe Int
go Int
lo (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
      where
        p :: Int
p = Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Integer -> Int
forall a. Num a => Integer -> a
fromInteger ((Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Array Int Integer
a Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
lo) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
hi Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lo) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` (Array Int Integer
a Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
hi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Array Int Integer
a Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
lo))

{- | Given a target and ascending finite list, grow a search range then find a matching
index or Nothing. Reject unsorted input. The implementation returns a matching index
without a general first-duplicate guarantee.
-}
exponentialSearch :: (Ord a) => a -> [a] -> Either MathError (Maybe Int)
exponentialSearch :: forall a. Ord a => a -> [a] -> Either MathError (Maybe Int)
exponentialSearch a
x [a]
xs = do
    [a] -> Either MathError ()
forall a. Ord a => [a] -> Either MathError ()
ordered [a]
xs
    let n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
        hi :: Int
hi = Int -> Int -> Int
grow Int
1 Int
n
        lo :: Int
lo = Int
hi Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
        i :: Int
i = Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (a -> a -> Bool) -> a -> [a] -> Int
forall a. Ord a => (a -> a -> Bool) -> a -> [a] -> Int
bound a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<) a
x (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take (Int
hi Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
lo [a]
xs))
    Maybe Int -> Either MathError (Maybe Int)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n Bool -> Bool -> Bool
&& [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x then Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i else Maybe Int
forall a. Maybe a
Nothing)
  where
    grow :: Int -> Int -> Int
grow Int
k Int
n | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n Bool -> Bool -> Bool
|| [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
k a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
x = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
k (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) | Bool
otherwise = Int -> Int -> Int
grow (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k) Int
n

{- | Fixed-iteration minimization of a unimodal finite function on [lo,hi].
Returns the final bracket; unimodality is a caller precondition.
-}
ternarySearch :: Int -> (Double -> Double) -> Double -> Double -> Either MathError (Double, Double)
ternarySearch :: Int
-> (Double -> Double)
-> Double
-> Double
-> Either MathError (Double, Double)
ternarySearch Int
steps Double -> Double
f Double
lo Double
hi
    | Int
steps Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
steps Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
100000 Bool -> Bool -> Bool
|| Bool -> Bool
not ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
forall {a}. RealFloat a => a -> Bool
finite [Double
lo, Double
hi]) Bool -> Bool -> Bool
|| Double
lo Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
hi =
        MathError -> Either MathError (Double, Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid ternary search interval/steps")
    | Bool
otherwise = Int -> Double -> Double -> Either MathError (Double, Double)
forall {t}.
(Eq t, Num t) =>
t -> Double -> Double -> Either MathError (Double, Double)
go Int
steps Double
lo Double
hi
  where
    finite :: a -> Bool
finite a
x = Bool -> Bool
not (a -> Bool
forall {a}. RealFloat a => a -> Bool
isNaN a
x Bool -> Bool -> Bool
|| a -> Bool
forall {a}. RealFloat a => a -> Bool
isInfinite a
x)
    go :: t -> Double -> Double -> Either MathError (Double, Double)
go t
0 Double
a Double
b = (Double, Double) -> Either MathError (Double, Double)
forall a b. b -> Either a b
Right (Double
a, Double
b)
    go t
k Double
a Double
b =
        let c :: Double
c = Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
3; d :: Double
d = Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
3; u :: Double
u = Double -> Double
f Double
c; v :: Double
v = Double -> Double
f Double
d
         in if Bool -> Bool
not ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
forall {a}. RealFloat a => a -> Bool
finite [Double
c, Double
d, Double
u, Double
v])
                then MathError -> Either MathError (Double, Double)
forall a b. a -> Either a b
Left MathError
NonFiniteResult
                else if Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
v then t -> Double -> Double -> Either MathError (Double, Double)
go (t
k t -> t -> t
forall a. Num a => a -> a -> a
- t
1) Double
a Double
d else t -> Double -> Double -> Either MathError (Double, Double)
go (t
k t -> t -> t
forall a. Num a => a -> a -> a
- t
1) Double
c Double
b

select :: (Ord a) => Int -> Int -> [a] -> a
select :: forall a. Ord a => Int -> Int -> [a] -> a
select Int
depth Int
k [a]
xs
    | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
0
    | Int
depth Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = [a] -> [a]
forall a. Ord a => [a] -> [a]
heapsort [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
k
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
lo = Int -> Int -> [a] -> a
forall a. Ord a => Int -> Int -> [a] -> a
select (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
k [a]
lo
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
eq = a
p
    | Bool
otherwise = Int -> Int -> [a] -> a
forall a. Ord a => Int -> Int -> [a] -> a
select (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
- [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
eq) [a]
hi
  where
    p :: a
p = [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2); lo :: [a]
lo = (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
p) [a]
xs; eq :: [a]
eq = (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
p) [a]
xs; hi :: [a]
hi = (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
p) [a]
xs
quickselect, introselect :: (Ord a) => Int -> [a] -> Either MathError a

{- | Given a zero-based order-statistic index and list, return the element at that position
in ascending order without returning a sorted list. Reject negative/out-of-range
indices. Example: quickselect 1 [9,1,5] == Right 5.
-}
quickselect :: forall a. Ord a => Int -> [a] -> Either MathError a
quickselect Int
k [a]
xs
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs = MathError -> Either MathError a
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"selection index out of bounds")
    | Bool
otherwise = a -> Either MathError a
forall a b. b -> Either a b
Right (Int -> Int -> [a] -> a
forall a. Ord a => Int -> Int -> [a] -> a
select ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) Int
k [a]
xs)

{- | Like quickselect, but bound pivot depth and fall back to heap sorting. Index must lie
in [0,length). Example: introselect 0 [3,2,1] == Right 1.
-}
introselect :: forall a. Ord a => Int -> [a] -> Either MathError a
introselect Int
k [a]
xs
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs = MathError -> Either MathError a
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"selection index out of bounds")
    | Bool
otherwise = a -> Either MathError a
forall a b. b -> Either a b
Right (Int -> Int -> [a] -> a
forall a. Ord a => Int -> Int -> [a] -> a
select (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int
log2 ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)) Int
k [a]
xs)

log2 :: Int -> Int
log2 :: Int -> Int
log2 Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = Int
0 | Bool
otherwise = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
log2 (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)

{- | Given k in [0,length] and a list, return the k smallest values in ascending order,
retaining duplicates. Invalid k fails. Example: partialSort 2 [5,1,3] == Right [1,3].
-}
partialSort :: (Ord a) => Int -> [a] -> Either MathError [a]
partialSort :: forall a. Ord a => Int -> [a] -> Either MathError [a]
partialSort Int
k [a]
xs
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs = MathError -> Either MathError [a]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"partial-sort size out of bounds")
    | Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = [a] -> Either MathError [a]
forall a b. b -> Either a b
Right []
    | Bool
otherwise = do
        p <- Int -> [a] -> Either MathError a
forall a. Ord a => Int -> [a] -> Either MathError a
introselect (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
xs
        let lo = (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
p) [a]
xs
        pure (mergeSort lo ++ take (k - length lo) (filter (== p) xs))

{- | Stable ascending merge sort under a lawful total Ord; accepts empty lists. Return a new
sorted list. Example: mergeSort [3,1,2] == [1,2,3].
-}
mergeSort :: (Ord a) => [a] -> [a]
mergeSort :: forall a. Ord a => [a] -> [a]
mergeSort [] = []
mergeSort [a
x] = [a
x]
mergeSort [a]
xs = [a] -> [a] -> [a]
forall {a}. Ord a => [a] -> [a] -> [a]
merge ([a] -> [a]
forall a. Ord a => [a] -> [a]
mergeSort [a]
a) ([a] -> [a]
forall a. Ord a => [a] -> [a]
mergeSort [a]
b)
  where
    ([a]
a, [a]
b) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [a]
xs
    merge :: [a] -> [a] -> [a]
merge [] [a]
ys = [a]
ys
    merge [a]
ys [] = [a]
ys
    merge aa :: [a]
aa@(a
u : [a]
us) bb :: [a]
bb@(a
v : [a]
vs) | a
u a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
v = a
u a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
merge [a]
us [a]
bb | Bool
otherwise = a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
merge [a]
aa [a]
vs

{- | Ascending reference quicksort preserving equal values, using the head as pivot.
Worst-case quadratic work; accepts empty lists. Example: quicksort [2,1,2] == [1,2,2].
-}
quicksort :: (Ord a) => [a] -> [a]
quicksort :: forall a. Ord a => [a] -> [a]
quicksort [] = []
quicksort (a
p : [a]
xs) = [a] -> [a]
forall a. Ord a => [a] -> [a]
quicksort ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
p) [a]
xs) [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ (a
p a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
p) [a]
xs) [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a] -> [a]
forall a. Ord a => [a] -> [a]
quicksort ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
p) [a]
xs)

{- | In-place-style array max heap with persistent updates. Correct heap sort;
persistence adds allocation to the usual O(n log n) comparisons.
-}
heapsort :: (Ord a) => [a] -> [a]
heapsort :: forall a. Ord a => [a] -> [a]
heapsort [a]
xs = Array Int a -> [a]
forall i e. Array i e -> [e]
elems (Int -> Array Int a -> Array Int a
forall {t} {e}. (Num t, Ord e, Ix t) => t -> Array t e -> Array t e
extract (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Array Int a
built)
  where
    n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
    a :: Array Int a
a = (Int, Int) -> [a] -> Array Int a
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
xs
    swap :: Array i e -> i -> i -> Array i e
swap Array i e
ar i
i i
j = Array i e
ar Array i e -> [(i, e)] -> Array i e
forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(i
i, Array i e
ar Array i e -> i -> e
forall i e. Ix i => Array i e -> i -> e
! i
j), (i
j, Array i e
ar Array i e -> i -> e
forall i e. Ix i => Array i e -> i -> e
! i
i)]
    sift :: Array t e -> t -> t -> Array t e
sift Array t e
ar t
size t
i
        | t
left t -> t -> Bool
forall a. Ord a => a -> a -> Bool
>= t
size = Array t e
ar
        | Array t e
ar Array t e -> t -> e
forall i e. Ix i => Array i e -> i -> e
! t
i e -> e -> Bool
forall a. Ord a => a -> a -> Bool
>= Array t e
ar Array t e -> t -> e
forall i e. Ix i => Array i e -> i -> e
! t
child = Array t e
ar
        | Bool
otherwise = Array t e -> t -> t -> Array t e
sift (Array t e -> t -> t -> Array t e
forall {i} {e}. Ix i => Array i e -> i -> i -> Array i e
swap Array t e
ar t
i t
child) t
size t
child
      where
        left :: t
left = t
2 t -> t -> t
forall a. Num a => a -> a -> a
* t
i t -> t -> t
forall a. Num a => a -> a -> a
+ t
1
        right :: t
right = t
left t -> t -> t
forall a. Num a => a -> a -> a
+ t
1
        child :: t
child = if t
right t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< t
size Bool -> Bool -> Bool
&& Array t e
ar Array t e -> t -> e
forall i e. Ix i => Array i e -> i -> e
! t
right e -> e -> Bool
forall a. Ord a => a -> a -> Bool
> Array t e
ar Array t e -> t -> e
forall i e. Ix i => Array i e -> i -> e
! t
left then t
right else t
left
    built :: Array Int a
built = (Array Int a -> Int -> Array Int a)
-> Array Int a -> [Int] -> Array Int a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\Array Int a
ar Int
i -> Array Int a -> Int -> Int -> Array Int a
forall {t} {e}.
(Num t, Ord e, Ix t) =>
Array t e -> t -> t -> Array t e
sift Array Int a
ar Int
n Int
i) Array Int a
a [Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
0]
    extract :: t -> Array t e -> Array t e
extract t
i Array t e
ar | t
i t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 = Array t e
ar | Bool
otherwise = t -> Array t e -> Array t e
extract (t
i t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (Array t e -> t -> t -> Array t e
forall {t} {e}.
(Num t, Ord e, Ix t) =>
Array t e -> t -> t -> Array t e
sift (Array t e -> t -> t -> Array t e
forall {i} {e}. Ix i => Array i e -> i -> i -> Array i e
swap Array t e
ar t
0 t
i) t
i t
0)

{- | Ascending comparison sort with bounded quickselect-style partition depth and heap-sort
fallback. Return a new sorted list; requires a lawful total Ord. Example: introsort
[3,1,2] == [1,2,3].
-}
introsort :: (Ord a) => [a] -> [a]
introsort :: forall a. Ord a => [a] -> [a]
introsort [a]
xs = Int -> [a] -> [a]
forall {t} {a}. (Num t, Ord a, Eq t) => t -> [a] -> [a]
go (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int
log2 ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)) [a]
xs
  where
    go :: t -> [a] -> [a]
go t
_ [] = []
    go t
0 [a]
ys = [a] -> [a]
forall a. Ord a => [a] -> [a]
heapsort [a]
ys
    go t
d [a]
ys =
        let p :: a
p = [a]
ys [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
         in t -> [a] -> [a]
go (t
d t -> t -> t
forall a. Num a => a -> a -> a
- t
1) ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
p) [a]
ys) [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
p) [a]
ys [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ t -> [a] -> [a]
go (t
d t -> t -> t
forall a. Num a => a -> a -> a
- t
1) ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
p) [a]
ys)

{- | Given a nonnegative value-range budget then an Int list, return ascending counting
sort. Reject a negative budget or a range max-min+1 exceeding it (NonConvergence);
empty input succeeds. Example: countingSort 5 [3,1,3] == Right [1,3,3].
-}
countingSort :: Int -> [Int] -> Either MathError [Int]
countingSort :: Int -> [Int] -> Either MathError [Int]
countingSort Int
budget [Int]
xs
    | Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError [Int]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative counting-sort range budget")
    | [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
xs = [Int] -> Either MathError [Int]
forall a b. b -> Either a b
Right []
    | Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
hi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
lo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
budget =
        MathError -> Either MathError [Int]
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"counting sort range" Int
budget)
    | Bool
otherwise = [Int] -> Either MathError [Int]
forall a b. b -> Either a b
Right ([[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (Int -> Int -> Map Int Int -> Int
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Int
0 Int
x Map Int Int
counts) Int
x | Int
x <- [Int
lo .. Int
hi]])
  where
    lo :: Int
lo = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Int]
xs; hi :: Int
hi = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Int]
xs; counts :: Map Int Int
counts = (Int -> Int -> Int) -> [(Int, Int)] -> Map Int Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [(Int
x, Int
1) | Int
x <- [Int]
xs]

{- | Sort unbounded Integers in ascending order by radix digits, separating negative and
nonnegative magnitudes. Accepts empty input. Example: radixSort [2,-3,0] == [-3,0,2].
-}
radixSort :: [Integer] -> [Integer]
radixSort :: [Integer] -> [Integer]
radixSort [] = []
radixSort [Integer]
xs = (Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
lo) (Integer -> [Integer] -> [Integer]
go Integer
1 [Integer]
shifted)
  where
    lo :: Integer
lo = [Integer] -> Integer
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Integer]
xs
    shifted :: [Integer]
shifted = (Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
subtract Integer
lo) [Integer]
xs
    hi :: Integer
hi = [Integer] -> Integer
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Integer]
shifted
    go :: Integer -> [Integer] -> [Integer]
go Integer
p [Integer]
ys
        | Integer
p Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
hi = [Integer]
ys
        | Bool
otherwise = Integer -> [Integer] -> [Integer]
go (Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
256) ([[Integer]] -> [Integer]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Integer
x | Integer
x <- [Integer]
ys, (Integer
x Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
p) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
256 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
d] | Integer
d <- [Integer
0 .. Integer
255]])

{- | Given bucket count in [1,65536] and finite Doubles in [0,1], sort ascending using
buckets and merge sort. Invalid count/data fail. Example: bucketSort 4 [1,0.5,0] ==
Right [0,0.5,1].
-}
bucketSort :: Int -> [Double] -> Either MathError [Double]
bucketSort :: Int -> [Double] -> Either MathError [Double]
bucketSort Int
buckets [Double]
xs
    | Int
buckets Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int
buckets Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
65536 = MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"bucket count must be in [1,65536]")
    | (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Double
x -> Double -> Bool
forall {a}. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall {a}. RealFloat a => a -> Bool
isInfinite Double
x Bool -> Bool -> Bool
|| Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1) [Double]
xs =
        MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"bucket sort values must be in [0,1]")
    | Bool
otherwise = [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double] -> [Double]
forall a. Ord a => [a] -> [a]
mergeSort [Double
x | Double
x <- [Double]
xs, Double -> Int
forall {a}. RealFrac a => a -> Int
index Double
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i] | Int
i <- [Int
0 .. Int
buckets Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
  where
    index :: a -> Int
index a
x = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
buckets Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (a -> Int
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (a
x a -> a -> a
forall a. Num a => a -> a -> a
* Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
buckets))

{- | Return original zero-based indices in stable ascending value order. Example: argsort
[3,1,1] == [1,2,0]. No values are discarded.
-}
argsort :: (Ord a) => [a] -> [Int]
argsort :: forall a. Ord a => [a] -> [Int]
argsort = ((a, Int) -> Int) -> [(a, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (a, Int) -> Int
forall a b. (a, b) -> b
snd ([(a, Int)] -> [Int]) -> ([a] -> [(a, Int)]) -> [a] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, Int)] -> [(a, Int)]
forall a. Ord a => [a] -> [a]
mergeSort ([(a, Int)] -> [(a, Int)])
-> ([a] -> [(a, Int)]) -> [a] -> [(a, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> [Int] -> [(a, Int)]) -> [Int] -> [a] -> [(a, Int)]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [a] -> [Int] -> [(a, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..]

{- | Alias of partialSort: return the k smallest values in ascending order, not the k
largest. k outside [0,length] fails. Example: topK 2 [5,1,3] == Right [1,3].
-}
topK :: (Ord a) => Int -> [a] -> Either MathError [a]
topK :: forall a. Ord a => Int -> [a] -> Either MathError [a]
topK = Int -> [a] -> Either MathError [a]
forall a. Ord a => Int -> [a] -> Either MathError [a]
partialSort

{- | Given finite q in [0,1] and a nonempty finite Double list, linearly interpolate the
ascending order statistics at q*(n-1). Invalid data fail. Example: quantileSelect 0.5
[1,3] == Right 2.
-}
quantileSelect :: Double -> [Double] -> Either MathError Double
quantileSelect :: Double -> [Double] -> Either MathError Double
quantileSelect Double
q [Double]
xs
    | [Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs Bool -> Bool -> Bool
|| Double
q Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Double
q Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1 Bool -> Bool -> Bool
|| (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Double
x -> Double -> Bool
forall {a}. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall {a}. RealFloat a => a -> Bool
isInfinite Double
x) (Double
q Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
xs) =
        MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid quantile input")
    | Bool
otherwise = do a <- Int -> [Double] -> Either MathError Double
forall a. Ord a => Int -> [a] -> Either MathError a
introselect Int
lo [Double]
xs; b <- introselect hi xs; pure ((1 - t) * a + t * b)
  where
    p :: Double
p = Double
q Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1); lo :: Int
lo = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
p; hi :: Int
hi = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling Double
p; t :: Double
t = Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
lo

{- | Sort ascending and keep one representative of each equal value. Input order is not
preserved. Example: deduplicate [3,1,3] == [1,3].
-}
deduplicate :: (Ord a) => [a] -> [a]
deduplicate :: forall a. Ord a => [a] -> [a]
deduplicate [a]
xs = [a
x | a
x : [a]
_ <- [a] -> [[a]]
forall a. Eq a => [a] -> [[a]]
group ([a] -> [a]
forall a. Ord a => [a] -> [a]
mergeSort [a]
xs)]

{- | Sort ascending and return each distinct value paired with its frequency. Example:
numericGroups [2,1,2] == [(1,1),(2,2)].
-}
numericGroups :: (Ord a) => [a] -> [(a, Int)]
numericGroups :: forall a. Ord a => [a] -> [(a, Int)]
numericGroups [a]
xs = [(a
x, [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
g) | g :: [a]
g@(a
x : [a]
_) <- [a] -> [[a]]
forall a. Eq a => [a] -> [[a]]
group ([a] -> [a]
forall a. Ord a => [a] -> [a]
mergeSort [a]
xs)]

{- | Encode adjacent equal runs as (value,count), preserving run order without sorting.
Example: runLengthEncode [1,1,2,1] == [(1,2),(2,1),(1,1)].
-}
runLengthEncode :: (Eq a) => [a] -> [(a, Int)]
runLengthEncode :: forall a. Eq a => [a] -> [(a, Int)]
runLengthEncode [a]
xs = [(a
x, [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
g) | g :: [a]
g@(a
x : [a]
_) <- [a] -> [[a]]
forall a. Eq a => [a] -> [[a]]
group [a]
xs]

{- | Given maximum output length then (value,count) runs, expand runs. Negative budget/count
fails; output beyond the budget returns NonConvergence before allocation. Zero-length
runs are allowed.
-}
runLengthDecode :: Int -> [(a, Int)] -> Either MathError [a]
runLengthDecode :: forall a. Int -> [(a, Int)] -> Either MathError [a]
runLengthDecode Int
b [(a, Int)]
runs
    | Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| ((a, Int) -> Bool) -> [(a, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (Int -> Bool) -> ((a, Int) -> Int) -> (a, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Int) -> Int
forall a b. (a, b) -> b
snd) [(a, Int)]
runs = MathError -> Either MathError [a]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative RLE size")
    | [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((a, Int) -> Integer) -> [(a, Int)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Integer) -> ((a, Int) -> Int) -> (a, Int) -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Int) -> Int
forall a b. (a, b) -> b
snd) [(a, Int)]
runs) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
b = MathError -> Either MathError [a]
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"RLE output" Int
b)
    | Bool
otherwise = [a] -> Either MathError [a]
forall a b. b -> Either a b
Right (((a, Int) -> [a]) -> [(a, Int)] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((a -> Int -> [a]) -> (a, Int) -> [a]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Int -> a -> [a]) -> a -> Int -> [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> a -> [a]
forall a. Int -> a -> [a]
replicate)) [(a, Int)]
runs)

prefixSums, suffixSums, differenceArray, applyDifferences :: (Num a) => [a] -> [a]

{- | Return cumulative sums including a leading zero; output length is n+1. Arithmetic
follows Num. Example: prefixSums [1,2] == [0,1,3].
-}
prefixSums :: forall a. Num a => [a] -> [a]
prefixSums = (a -> a -> a) -> a -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0

{- | Return suffix sums including a final zero; output length is n+1. Example: suffixSums
[1,2] == [3,2,0].
-}
suffixSums :: forall a. Num a => [a] -> [a]
suffixSums = (a -> a -> a) -> a -> [a] -> [a]
forall a b. (a -> b -> b) -> b -> [a] -> [b]
scanr a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0

{- | Return first differences relative to a leading zero, preserving length. Empty input
gives empty output. Example: differenceArray [2,5,9] == [2,3,4].
-}
differenceArray :: forall a. Num a => [a] -> [a]
differenceArray [] = []
differenceArray [a]
xs = (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [a]
xs (a
0 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs)

{- | Reconstruct values from first differences by inclusive cumulative sum. Example:
applyDifferences [2,3,4] == [2,5,9].
-}
applyDifferences :: forall a. Num a => [a] -> [a]
applyDifferences = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
1 ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. Num a => [a] -> [a]
prefixSums

{- | Return sorted unique keys and zero-based ranks for the original input order. Example:
coordinateCompression [8,2,8] == ([2,8],[1,0,1]).
-}
coordinateCompression :: (Ord a) => [a] -> ([a], [Int])
coordinateCompression :: forall a. Ord a => [a] -> ([a], [Int])
coordinateCompression [a]
xs = ([a]
keys, (a -> Int) -> [a] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> Map a Int
table Map a Int -> a -> Int
forall k a. Ord k => Map k a -> k -> a
M.! a
x) [a]
xs)
  where
    keys :: [a]
keys = [a] -> [a]
forall a. Ord a => [a] -> [a]
deduplicate [a]
xs; table :: Map a Int
table = [(a, Int)] -> Map a Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([a] -> [Int] -> [(a, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
keys [Int
0 ..])

{- | Half-open interval union length and peak overlap; zero-length intervals
contribute neither. End events precede starts at the same coordinate.
-}
intervalSweep :: [(Integer, Integer)] -> Either MathError (Integer, Int)
intervalSweep :: [(Integer, Integer)] -> Either MathError (Integer, Int)
intervalSweep [(Integer, Integer)]
intervals
    | ((Integer, Integer) -> Bool) -> [(Integer, Integer)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Integer -> Integer -> Bool) -> (Integer, Integer) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
(>)) [(Integer, Integer)]
intervals = MathError -> Either MathError (Integer, Int)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"reversed interval")
    | Bool
otherwise = (Integer, Int) -> Either MathError (Integer, Int)
forall a b. b -> Either a b
Right ([(Integer, Int)]
-> Integer -> Int -> Integer -> Int -> (Integer, Int)
forall {t} {t}.
(Num t, Num t, Ord t) =>
[(t, t)] -> t -> t -> t -> t -> (t, t)
go [(Integer, Int)]
events Integer
0 Int
0 Integer
0 Int
0)
  where
    events :: [(Integer, Int)]
events = [(Integer, Int)] -> [(Integer, Int)]
forall a. Ord a => [a] -> [a]
sort ([[(Integer, Int)]] -> [(Integer, Int)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(Integer
a, Int
1 :: Int), (Integer
b, -Int
1)] | (Integer
a, Integer
b) <- [(Integer, Integer)]
intervals, Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
b])
    go :: [(t, t)] -> t -> t -> t -> t -> (t, t)
go [] t
_ t
_ t
total t
peak = (t
total, t
peak)
    go ((t
x, t
d) : [(t, t)]
rest) t
old t
active t
total t
peak = [(t, t)] -> t -> t -> t -> t -> (t, t)
go [(t, t)]
rest t
x (t
active t -> t -> t
forall a. Num a => a -> a -> a
+ t
d) (t
total t -> t -> t
forall a. Num a => a -> a -> a
+ if t
active t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> t
0 then t
x t -> t -> t
forall a. Num a => a -> a -> a
- t
old else t
0) (t -> t -> t
forall a. Ord a => a -> a -> a
max t
peak (t
active t -> t -> t
forall a. Num a => a -> a -> a
+ t
d))

{- | One DP layer @next[j] = min_{0<=k<j} previous[k]+cost k j@.
Monotone argmins are checked against all candidates before returning;
invalid optimization assumptions fail explicitly. O(n^2) validation.
-}
divideConquerDP :: [Double] -> (Int -> Int -> Double) -> Either MathError [Double]
divideConquerDP :: [Double] -> (Int -> Int -> Double) -> Either MathError [Double]
divideConquerDP [Double]
prev Int -> Int -> Double
cost
    | [Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
prev Bool -> Bool -> Bool
|| [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
prev Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
4096 Bool -> Bool -> Bool
|| (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Double -> Bool
forall {a}. RealFloat a => a -> Bool
bad [Double]
prev = MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid DP layer")
    | Bool
otherwise = do
        let n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
prev
            solve :: Int -> Int -> Int -> Int -> [(Int, Double, Int)]
solve Int
l Int
r Int
a Int
b
                | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
r = []
                | Bool
otherwise =
                    let m :: Int
m = (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
                        candidates :: [(Double, Int)]
candidates = [([Double]
prev [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Int -> Double
cost Int
j Int
m, Int
j) | Int
j <- [Int
a .. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
b (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]]
                        (Double
v, Int
k) = [(Double, Int)] -> (Double, Int)
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [(Double, Int)]
candidates
                     in Int -> Int -> Int -> Int -> [(Int, Double, Int)]
solve Int
l (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
a Int
k [(Int, Double, Int)]
-> [(Int, Double, Int)] -> [(Int, Double, Int)]
forall a. [a] -> [a] -> [a]
++ [(Int
m, Double
v, Int
k)] [(Int, Double, Int)]
-> [(Int, Double, Int)] -> [(Int, Double, Int)]
forall a. [a] -> [a] -> [a]
++ Int -> Int -> Int -> Int -> [(Int, Double, Int)]
solve (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
r Int
k Int
b
            out :: [(Int, Double, Int)]
out = ((Int, Double, Int) -> Int)
-> [(Int, Double, Int)] -> [(Int, Double, Int)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (\(Int
i, Double
_, Int
_) -> Int
i) (Int -> Int -> Int -> Int -> [(Int, Double, Int)]
solve Int
1 Int
n Int
0 (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
        if ((Int, Double, Int) -> Bool) -> [(Int, Double, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Int
j, Double
v, Int
_) -> Double -> Bool
forall {a}. RealFloat a => a -> Bool
bad Double
v Bool -> Bool -> Bool
|| Double
v Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
/= [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [[Double]
prev [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Int -> Double
cost Int
k Int
j | Int
k <- [Int
0 .. Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]) [(Int, Double, Int)]
out
            then MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"nonfinite costs or nonmonotone DP argmins")
            else [Double] -> Either MathError [Double]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Double
v | (Int
_, Double
v, Int
_) <- [(Int, Double, Int)]
out]
  where
    bad :: a -> Bool
bad a
x = a -> Bool
forall {a}. RealFloat a => a -> Bool
isNaN a
x Bool -> Bool -> Bool
|| a -> Bool
forall {a}. RealFloat a => a -> Bool
isInfinite a
x

-- | Bottom-up bounded DAG recurrence; dependencies must have lower indices.
tabulateDP :: Int -> (Int -> [(Int, Integer)]) -> Either MathError [Integer]
tabulateDP :: Int -> (Int -> [(Int, Integer)]) -> Either MathError [Integer]
tabulateDP Int
n Int -> [(Int, Integer)]
recurrence
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
100000 = MathError -> Either MathError [Integer]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"DP size exceeds 100000")
    | Bool
otherwise = Int -> [Integer] -> Either MathError [Integer]
go Int
0 []
  where
    go :: Int -> [Integer] -> Either MathError [Integer]
go Int
i [Integer]
values
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = [Integer] -> Either MathError [Integer]
forall a b. b -> Either a b
Right [Integer]
values
        | Bool
otherwise =
            let terms :: [(Int, Integer)]
terms = Int -> [(Int, Integer)]
recurrence Int
i
             in if ((Int, Integer) -> Bool) -> [(Int, Integer)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Int
j, Integer
_) -> Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
i) [(Int, Integer)]
terms
                    then MathError -> Either MathError [Integer]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"DP dependency is not earlier")
                    else Int -> [Integer] -> Either MathError [Integer]
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ([Integer]
values [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++ [if [(Int, Integer)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Int, Integer)]
terms then Integer
1 else [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer
c Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* [Integer]
values [Integer] -> Int -> Integer
forall a. HasCallStack => [a] -> Int -> a
!! Int
j | (Int
j, Integer
c) <- [(Int, Integer)]
terms]])

{- | Top-down memoization with a node budget. Recurrence supplies an additive
constant and weighted dependencies; cycles are rejected.
-}
memoizedDP :: Int -> (Int -> (Integer, [(Int, Integer)])) -> Int -> Either MathError Integer
memoizedDP :: Int
-> (Int -> (Integer, [(Int, Integer)]))
-> Int
-> Either MathError Integer
memoizedDP Int
b Int -> (Integer, [(Int, Integer)])
recurrence Int
target
    | Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative memoization budget")
    | Bool
otherwise = do (v, _, _) <- Int
-> [Int]
-> Map Int Integer
-> Int
-> Either MathError (Integer, Int, Map Int Integer)
forall {t}.
(Ord t, Num t) =>
t
-> [Int]
-> Map Int Integer
-> Int
-> Either MathError (Integer, t, Map Int Integer)
go Int
b [] Map Int Integer
forall k a. Map k a
M.empty Int
target; pure v
  where
    go :: t
-> [Int]
-> Map Int Integer
-> Int
-> Either MathError (Integer, t, Map Int Integer)
go t
fuel [Int]
stack Map Int Integer
cache Int
x = case Int -> Map Int Integer -> Maybe Integer
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
x Map Int Integer
cache of
        Just Integer
v -> (Integer, t, Map Int Integer)
-> Either MathError (Integer, t, Map Int Integer)
forall a b. b -> Either a b
Right (Integer
v, t
fuel, Map Int Integer
cache)
        Maybe Integer
Nothing
            | Int
x Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
stack -> MathError -> Either MathError (Integer, t, Map Int Integer)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"cyclic recurrence")
            | t
fuel t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 -> MathError -> Either MathError (Integer, t, Map Int Integer)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"memoized DP" Int
b)
            | Bool
otherwise -> do
                let (Integer
base, [(Int, Integer)]
terms) = Int -> (Integer, [(Int, Integer)])
recurrence Int
x
                (v, f, c) <-
                    (Either MathError (Integer, t, Map Int Integer)
 -> (Int, Integer)
 -> Either MathError (Integer, t, Map Int Integer))
-> Either MathError (Integer, t, Map Int Integer)
-> [(Int, Integer)]
-> Either MathError (Integer, t, Map Int Integer)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl'
                        ( \Either MathError (Integer, t, Map Int Integer)
result (Int
j, Integer
w) -> do (v, f, c) <- Either MathError (Integer, t, Map Int Integer)
result; (u, f', c') <- go f (x : stack) c j; pure (v + w * u, f', c')
                        )
                        ((Integer, t, Map Int Integer)
-> Either MathError (Integer, t, Map Int Integer)
forall a b. b -> Either a b
Right (Integer
base, t
fuel t -> t -> t
forall a. Num a => a -> a -> a
- t
1, Map Int Integer
cache))
                        [(Int, Integer)]
terms
                pure (v, f, M.insert x v c)

-- | K-fold application of a total finite successor map, by repeated squaring.
binaryLift :: [Int] -> Integer -> Int -> Either MathError Int
binaryLift :: [Int] -> Integer -> Int -> Either MathError Int
binaryLift [Int]
next Integer
steps Int
start
    | Integer
steps Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
next Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
next) [Int]
next =
        MathError -> Either MathError Int
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid lifting map/index")
    | Bool
otherwise = Int -> Either MathError Int
forall a b. b -> Either a b
Right ([Int] -> Integer -> Int -> Int
forall {a}. Integral a => [Int] -> a -> Int -> Int
go [Int]
next Integer
steps Int
start)
  where
    go :: [Int] -> a -> Int -> Int
go [Int]
_ a
0 Int
x = Int
x; go [Int]
table a
k Int
x = [Int] -> a -> Int -> Int
go ((Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([Int]
table [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) [Int]
table) (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2) (if a -> Bool
forall a. Integral a => a -> Bool
odd a
k then [Int]
table [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
x else Int
x)

{- | Given square Integer rows and a nonnegative exponent, return exact matrix
exponentiation by squaring; exponent zero returns identity. Reject ragged/non-square
matrices, order >128, or negative exponent.
-}
integerMatrixPower :: [[Integer]] -> Integer -> Either MathError [[Integer]]
integerMatrixPower :: [[Integer]] -> Integer -> Either MathError [[Integer]]
integerMatrixPower [[Integer]]
a Integer
powerIndex
    | Integer
powerIndex Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| [[Integer]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Integer]]
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
128 Bool -> Bool -> Bool
|| ([Integer] -> Bool) -> [[Integer]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n) (Int -> Bool) -> ([Integer] -> Int) -> [Integer] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Integer]]
a =
        MathError -> Either MathError [[Integer]]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid square integer matrix/exponent")
    | Bool
otherwise = [[Integer]] -> Either MathError [[Integer]]
forall a b. b -> Either a b
Right ([[Integer]] -> Integer -> [[Integer]] -> [[Integer]]
forall {a} {a}. (Num a, Integral a) => [[a]] -> a -> [[a]] -> [[a]]
go [[Integer]]
a Integer
powerIndex [[Integer]]
identity)
  where
    n :: Int
n = [[Integer]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Integer]]
a
    identity :: [[Integer]]
identity = [[if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j then Integer
1 else Integer
0 | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    mul :: [[a]] -> [[a]] -> [[a]]
mul [[a]]
x [[a]]
y = [[[a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[[a]]
x [[a]] -> Int -> [a]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
k a -> a -> a
forall a. Num a => a -> a -> a
* [[a]]
y [[a]] -> Int -> [a]
forall a. HasCallStack => [a] -> Int -> a
!! Int
k [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
j | Int
k <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    go :: [[a]] -> a -> [[a]] -> [[a]]
go [[a]]
_ a
0 [[a]]
r = [[a]]
r
    go [[a]]
x a
k [[a]]
r = [[a]] -> a -> [[a]] -> [[a]]
go ([[a]] -> [[a]] -> [[a]]
forall {a}. Num a => [[a]] -> [[a]] -> [[a]]
mul [[a]]
x [[a]]
x) (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2) (if a -> Bool
forall a. Integral a => a -> Bool
odd a
k then [[a]] -> [[a]] -> [[a]]
forall {a}. Num a => [[a]] -> [[a]] -> [[a]]
mul [[a]]
r [[a]]
x else [[a]]
r)

{- | Given base, modulus >1, then Integer sequence, fold h=(h*base+x) mod modulus from zero.
Invalid modulus fails. This is noncryptographic and collisions are possible.
-}
polynomialHash :: Integer -> Integer -> [Integer] -> Either MathError Integer
polynomialHash :: Integer -> Integer -> [Integer] -> Either MathError Integer
polynomialHash Integer
base Integer
modulus [Integer]
xs
    | Integer
modulus Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"hash modulus must exceed one")
    | Bool
otherwise = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right ((Integer -> Integer -> Integer) -> Integer -> [Integer] -> Integer
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\Integer
h Integer
x -> (Integer
h Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
base Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
x) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
modulus) Integer
0 [Integer]
xs)

-- | Noncryptographic polynomial rolling hash; equal hashes can collide.
rollingHashes :: Integer -> Integer -> Int -> [Integer] -> Either MathError [Integer]
rollingHashes :: Integer
-> Integer -> Int -> [Integer] -> Either MathError [Integer]
rollingHashes Integer
base Integer
modulus Int
width [Integer]
xs
    | Integer
modulus Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1 Bool -> Bool -> Bool
|| Int
width Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
width Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
xs =
        MathError -> Either MathError [Integer]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid rolling hash window/modulus")
    | Int
width Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = [Integer] -> Either MathError [Integer]
forall a b. b -> Either a b
Right (Int -> Integer -> [Integer]
forall a. Int -> a -> [a]
replicate ([Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Integer
0)
    | Bool
otherwise = do
        h <- Integer -> Integer -> [Integer] -> Either MathError Integer
polynomialHash Integer
base Integer
modulus (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
width [Integer]
xs)
        let power = Integer -> Int -> Integer -> Integer
forall {a}. Integral a => Integer -> a -> Integer -> Integer
pow Integer
base (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Integer
1
            go Integer
_ [] [Integer]
_ = []
            go Integer
_ [Integer]
_ [] = []
            go Integer
old (Integer
a : [Integer]
as) (Integer
b : [Integer]
bs) = let v :: Integer
v = ((Integer
old Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
power) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
base Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
b) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
modulus in Integer
v Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: Integer -> [Integer] -> [Integer] -> [Integer]
go Integer
v [Integer]
as [Integer]
bs
        pure (h : go h xs (drop width xs))
  where
    pow :: Integer -> a -> Integer -> Integer
pow Integer
_ a
0 Integer
r = Integer
r
    pow Integer
x a
k Integer
r = Integer -> a -> Integer -> Integer
pow (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
x Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
modulus) (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2) (if a -> Bool
forall a. Integral a => a -> Bool
odd a
k then Integer
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
x Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
modulus else Integer
r)