Speed up deleteLargestM

src/Data/List/Utils.hs

+{-# 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 #-}