{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
module Data.List.Extras.Argmax
(
catchNull
, argmaxBy, argmaxesBy, argmaxWithMaxBy, argmaxesWithMaxBy
, argmax, argmaxes, argmaxWithMax, argmaxesWithMax
, argmin, argmins, argminWithMin, argminsWithMin
) where
import Data.List (foldl')
catchNull :: ([a] -> b) -> ([a] -> Maybe b)
{-# INLINE catchNull #-}
catchNull :: ([a] -> b) -> [a] -> Maybe b
catchNull [a] -> b
f = \[a]
xs ->
case [a]
xs of
[] -> Maybe b
forall a. Maybe a
Nothing
a
_:[a]
_ -> b -> Maybe b
forall a. a -> Maybe a
Just ([a] -> b
f [a]
xs)
emptyListError :: String -> a
{-# NOINLINE emptyListError #-}
emptyListError :: String -> a
emptyListError String
fun =
String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"Data.List.Extras.Argmax."String -> String -> String
forall a. [a] -> [a] -> [a]
++String
funString -> String -> String
forall a. [a] -> [a] -> [a]
++String
": empty list"
throwNull :: String -> (a -> [a] -> b) -> ([a] -> b)
{-# INLINE throwNull #-}
throwNull :: String -> (a -> [a] -> b) -> [a] -> b
throwNull String
fun a -> [a] -> b
f = \[a]
xs ->
case [a]
xs of
[] -> String -> b
forall a. String -> a
emptyListError String
fun
a
x:[a]
xs' -> a -> [a] -> b
f a
x [a]
xs'
_argmaxWithMaxBy :: (b -> b -> Bool) -> (a -> b) -> a -> [a] -> (a,b)
{-# INLINE _argmaxWithMaxBy #-}
_argmaxWithMaxBy :: (b -> b -> Bool) -> (a -> b) -> a -> [a] -> (a, b)
_argmaxWithMaxBy b -> b -> Bool
isBetterThan a -> b
f =
\a
x [a]
xs -> ((a, b) -> a -> (a, b)) -> (a, b) -> [a] -> (a, b)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (a, b) -> a -> (a, b)
cmp (a
x, a -> b
f a
x) [a]
xs
where
cmp :: (a, b) -> a -> (a, b)
cmp bfb :: (a, b)
bfb@(a
_,b
fb) a
a =
let fa :: b
fa = a -> b
f a
a in
if b
fa b -> b -> Bool
`isBetterThan` b
fb
then (a
a,b
fa)
else (a, b)
bfb
_argmaxesWithMaxBy :: (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> ([a],b)
{-# INLINE _argmaxesWithMaxBy #-}
_argmaxesWithMaxBy :: (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> ([a], b)
_argmaxesWithMaxBy b -> b -> Ordering
isBetterEqualThan a -> b
f =
\a
x [a]
xs -> (([a], b) -> a -> ([a], b)) -> ([a], b) -> [a] -> ([a], b)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ([a], b) -> a -> ([a], b)
cmp ([a
x], a -> b
f a
x) [a]
xs
where
cmp :: ([a], b) -> a -> ([a], b)
cmp bsfb :: ([a], b)
bsfb@([a]
bs,b
fb) a
a =
let fa :: b
fa = a -> b
f a
a in
case b -> b -> Ordering
isBetterEqualThan b
fa b
fb of
Ordering
GT -> ([a
a], b
fa)
Ordering
EQ -> (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
bs, b
fb)
Ordering
_ -> ([a], b)
bsfb
_argmaxBy :: (b -> b -> Bool) -> (a -> b) -> a -> [a] -> a
{-# INLINE _argmaxBy #-}
_argmaxBy :: (b -> b -> Bool) -> (a -> b) -> a -> [a] -> a
_argmaxBy b -> b -> Bool
k a -> b
f = ((a, b) -> a
forall a b. (a, b) -> a
fst ((a, b) -> a) -> ([a] -> (a, b)) -> [a] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) (([a] -> (a, b)) -> [a] -> a)
-> (a -> [a] -> (a, b)) -> a -> [a] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> b -> Bool) -> (a -> b) -> a -> [a] -> (a, b)
forall b a. (b -> b -> Bool) -> (a -> b) -> a -> [a] -> (a, b)
_argmaxWithMaxBy b -> b -> Bool
k a -> b
f
_argmaxesBy :: (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> [a]
{-# INLINE _argmaxesBy #-}
_argmaxesBy :: (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> [a]
_argmaxesBy b -> b -> Ordering
k a -> b
f = (([a], b) -> [a]
forall a b. (a, b) -> a
fst (([a], b) -> [a]) -> ([a] -> ([a], b)) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) (([a] -> ([a], b)) -> [a] -> [a])
-> (a -> [a] -> ([a], b)) -> a -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> ([a], b)
forall b a.
(b -> b -> Ordering) -> (a -> b) -> a -> [a] -> ([a], b)
_argmaxesWithMaxBy b -> b -> Ordering
k a -> b
f
bool :: (a -> a -> Ordering) -> (a -> a -> Bool)
bool :: (a -> a -> Ordering) -> a -> a -> Bool
bool a -> a -> Ordering
ord = \a
a a
b -> a -> a -> Ordering
ord a
a a
b Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
GT
argmaxBy :: (b -> b -> Ordering) -> (a -> b) -> [a] -> a
argmaxBy :: (b -> b -> Ordering) -> (a -> b) -> [a] -> a
argmaxBy b -> b -> Ordering
ord a -> b
f = String -> (a -> [a] -> a) -> [a] -> a
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argmaxBy"
((a -> [a] -> a) -> [a] -> a) -> (a -> [a] -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$ (b -> b -> Bool) -> (a -> b) -> a -> [a] -> a
forall b a. (b -> b -> Bool) -> (a -> b) -> a -> [a] -> a
_argmaxBy ((b -> b -> Ordering) -> b -> b -> Bool
forall a. (a -> a -> Ordering) -> a -> a -> Bool
bool b -> b -> Ordering
ord) a -> b
f
argmaxesBy :: (b -> b -> Ordering) -> (a -> b) -> [a] -> [a]
argmaxesBy :: (b -> b -> Ordering) -> (a -> b) -> [a] -> [a]
argmaxesBy b -> b -> Ordering
ord a -> b
f = String -> (a -> [a] -> [a]) -> [a] -> [a]
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argmaxesBy"
((a -> [a] -> [a]) -> [a] -> [a])
-> (a -> [a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> [a]
forall b a. (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> [a]
_argmaxesBy b -> b -> Ordering
ord a -> b
f
argmaxWithMaxBy :: (b -> b -> Ordering) -> (a -> b) -> [a] -> (a, b)
argmaxWithMaxBy :: (b -> b -> Ordering) -> (a -> b) -> [a] -> (a, b)
argmaxWithMaxBy b -> b -> Ordering
ord a -> b
f = String -> (a -> [a] -> (a, b)) -> [a] -> (a, b)
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argmaxWithMaxBy"
((a -> [a] -> (a, b)) -> [a] -> (a, b))
-> (a -> [a] -> (a, b)) -> [a] -> (a, b)
forall a b. (a -> b) -> a -> b
$ (b -> b -> Bool) -> (a -> b) -> a -> [a] -> (a, b)
forall b a. (b -> b -> Bool) -> (a -> b) -> a -> [a] -> (a, b)
_argmaxWithMaxBy ((b -> b -> Ordering) -> b -> b -> Bool
forall a. (a -> a -> Ordering) -> a -> a -> Bool
bool b -> b -> Ordering
ord) a -> b
f
argmaxesWithMaxBy :: (b -> b -> Ordering) -> (a -> b) -> [a] -> ([a], b)
argmaxesWithMaxBy :: (b -> b -> Ordering) -> (a -> b) -> [a] -> ([a], b)
argmaxesWithMaxBy b -> b -> Ordering
ord a -> b
f = String -> (a -> [a] -> ([a], b)) -> [a] -> ([a], b)
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argmaxesWithMaxBy"
((a -> [a] -> ([a], b)) -> [a] -> ([a], b))
-> (a -> [a] -> ([a], b)) -> [a] -> ([a], b)
forall a b. (a -> b) -> a -> b
$ (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> ([a], b)
forall b a.
(b -> b -> Ordering) -> (a -> b) -> a -> [a] -> ([a], b)
_argmaxesWithMaxBy b -> b -> Ordering
ord a -> b
f
argmax :: (Ord b) => (a -> b) -> [a] -> a
argmax :: (a -> b) -> [a] -> a
argmax a -> b
f = String -> (a -> [a] -> a) -> [a] -> a
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argmax"
((a -> [a] -> a) -> [a] -> a) -> (a -> [a] -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$ (b -> b -> Bool) -> (a -> b) -> a -> [a] -> a
forall b a. (b -> b -> Bool) -> (a -> b) -> a -> [a] -> a
_argmaxBy b -> b -> Bool
forall a. Ord a => a -> a -> Bool
(>) a -> b
f
argmaxes :: (Ord b) => (a -> b) -> [a] -> [a]
argmaxes :: (a -> b) -> [a] -> [a]
argmaxes a -> b
f = String -> (a -> [a] -> [a]) -> [a] -> [a]
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argmaxes"
((a -> [a] -> [a]) -> [a] -> [a])
-> (a -> [a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> [a]
forall b a. (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> [a]
_argmaxesBy b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a -> b
f
argmaxWithMax :: (Ord b) => (a -> b) -> [a] -> (a, b)
argmaxWithMax :: (a -> b) -> [a] -> (a, b)
argmaxWithMax a -> b
f = String -> (a -> [a] -> (a, b)) -> [a] -> (a, b)
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argmaxWithMax"
((a -> [a] -> (a, b)) -> [a] -> (a, b))
-> (a -> [a] -> (a, b)) -> [a] -> (a, b)
forall a b. (a -> b) -> a -> b
$ (b -> b -> Bool) -> (a -> b) -> a -> [a] -> (a, b)
forall b a. (b -> b -> Bool) -> (a -> b) -> a -> [a] -> (a, b)
_argmaxWithMaxBy b -> b -> Bool
forall a. Ord a => a -> a -> Bool
(>) a -> b
f
argmaxesWithMax :: (Ord b) => (a -> b) -> [a] -> ([a], b)
argmaxesWithMax :: (a -> b) -> [a] -> ([a], b)
argmaxesWithMax a -> b
f = String -> (a -> [a] -> ([a], b)) -> [a] -> ([a], b)
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argmaxesWithMax"
((a -> [a] -> ([a], b)) -> [a] -> ([a], b))
-> (a -> [a] -> ([a], b)) -> [a] -> ([a], b)
forall a b. (a -> b) -> a -> b
$ (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> ([a], b)
forall b a.
(b -> b -> Ordering) -> (a -> b) -> a -> [a] -> ([a], b)
_argmaxesWithMaxBy b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a -> b
f
argmin :: (Ord b) => (a -> b) -> [a] -> a
argmin :: (a -> b) -> [a] -> a
argmin a -> b
f = String -> (a -> [a] -> a) -> [a] -> a
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argmax"
((a -> [a] -> a) -> [a] -> a) -> (a -> [a] -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$ (b -> b -> Bool) -> (a -> b) -> a -> [a] -> a
forall b a. (b -> b -> Bool) -> (a -> b) -> a -> [a] -> a
_argmaxBy b -> b -> Bool
forall a. Ord a => a -> a -> Bool
(<) a -> b
f
argmins :: (Ord b) => (a -> b) -> [a] -> [a]
argmins :: (a -> b) -> [a] -> [a]
argmins a -> b
f = String -> (a -> [a] -> [a]) -> [a] -> [a]
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argmins"
((a -> [a] -> [a]) -> [a] -> [a])
-> (a -> [a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> [a]
forall b a. (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> [a]
_argmaxesBy ((b -> b -> Ordering) -> b -> b -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
compare) a -> b
f
argminWithMin :: (Ord b) => (a -> b) -> [a] -> (a, b)
argminWithMin :: (a -> b) -> [a] -> (a, b)
argminWithMin a -> b
f = String -> (a -> [a] -> (a, b)) -> [a] -> (a, b)
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argminWithMin"
((a -> [a] -> (a, b)) -> [a] -> (a, b))
-> (a -> [a] -> (a, b)) -> [a] -> (a, b)
forall a b. (a -> b) -> a -> b
$ (b -> b -> Bool) -> (a -> b) -> a -> [a] -> (a, b)
forall b a. (b -> b -> Bool) -> (a -> b) -> a -> [a] -> (a, b)
_argmaxWithMaxBy b -> b -> Bool
forall a. Ord a => a -> a -> Bool
(<) a -> b
f
argminsWithMin :: (Ord b) => (a -> b) -> [a] -> ([a], b)
argminsWithMin :: (a -> b) -> [a] -> ([a], b)
argminsWithMin a -> b
f = String -> (a -> [a] -> ([a], b)) -> [a] -> ([a], b)
forall a b. String -> (a -> [a] -> b) -> [a] -> b
throwNull String
"argminsWithMin"
((a -> [a] -> ([a], b)) -> [a] -> ([a], b))
-> (a -> [a] -> ([a], b)) -> [a] -> ([a], b)
forall a b. (a -> b) -> a -> b
$ (b -> b -> Ordering) -> (a -> b) -> a -> [a] -> ([a], b)
forall b a.
(b -> b -> Ordering) -> (a -> b) -> a -> [a] -> ([a], b)
_argmaxesWithMaxBy ((b -> b -> Ordering) -> b -> b -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
compare) a -> b
f