diff --git a/src/Data/List/Utils.hs b/src/Data/List/Utils.hs index 31d3978113616e372d6c46922d4d260387b11e7a..267fe05dce24a53ebf41e1ff6c958816b7f09825 100644 --- a/src/Data/List/Utils.hs +++ b/src/Data/List/Utils.hs @@ -1,17 +1,24 @@ +{-# LANGUAGE BangPatterns #-} module Data.List.Utils ( deleteLargestM ) where -import Control.Monad.ST -import Data.Function (on) -import Data.List (delete,maximumBy) - -- | Given a size function and a list, delete one largest element from the list. -- --- This functions is monadic to allow the size function to be monadic. --- --- TODO This could probably be more efficient +-- This functions is monadic to allow the size function to be monadic. This size +-- function must only return positive integers. deleteLargestM :: (Eq e, Monad m) => (e -> m Int) -> [e] -> m [e] deleteLargestM sizeFunction lst = do - zipWithSize <- traverse (\x -> (,x) <$> sizeFunction x) lst - return (delete (snd (maximumBy (compare `on` fst) zipWithSize)) lst) + idx <- findMax 0 (-1) (-1) lst + if idx >= 0 + then return (take idx lst ++ drop (idx+1) lst) + else return lst + + where + findMax _ !maxIdx _ [] = return maxIdx + findMax !currentIdx !maxIdx !maxVal (x:xs) = do + v <- sizeFunction x + if v > maxVal + then findMax (currentIdx + 1) currentIdx v xs + else findMax (currentIdx + 1) maxIdx maxVal xs +{-# INLINE deleteLargestM #-}