Add special grouping functions

CHANGES.md

@@ -4,6 +4,8 @@ Unreleased
 1.7.3
 =====
 
+* [#255](https://github.com/serokell/universum/issues/255):
+  Add `groupByFst` and `groupByKey`.
 * [#236](https://github.com/serokell/universum/issues/236):
   Add `updateMVar'` and `updateTVar'`.
 * [#244](https://github.com/serokell/universum/issues/244)

src/Universum/Container/Utils.hs

@@ -0,0 +1,94 @@
+{-# LANGUAGE Trustworthy #-}
+
+-- | Utility functions on top of 'Container' typeclass.
+
+module Universum.Container.Utils
+       ( groupByFst
+       , groupByKey
+       , groupByKeyBy
+
+       , groupOn
+       ) where
+
+import Data.Function (id, on, (.))
+import Data.List.NonEmpty (NonEmpty (..), groupBy)
+
+import Universum.Base
+import Universum.Container.Class
+
+-- $setup
+-- >>> import Data.Function
+-- >>> import GHC.Num
+-- >>> import Data.Text (toLower)
+
+-- | Variation of 'groupByKey' that accepts the comparison function on keys
+-- explicitly. This function must define a (computable) equivalence relation.
+--
+-- Among multiple keys appearing in the same group it will pick the leftmost
+-- one as the representer of the equivalence class.
+--
+-- >>> groupByKeyBy (<=) id [(1, 10), (2, 11), (0, 12), (1, 13), (3, 14)]
+-- [(1,10 :| [11]),(0,12 :| [13,14])]
+groupByKeyBy
+  :: Container t
+  => (k -> k -> Bool) -> (Element t -> (k, v)) -> t -> [(k, NonEmpty v)]
+groupByKeyBy kcmp split = start . toList
+  where
+    start [] = []
+    start (a : as)
+      | (k, v) <- split a
+      , let (ys, zs) = go k as
+      = (k, v :| ys) : zs
+
+    go _ [] = ([], [])
+    go ko (a : as)
+      | (kn, v) <- split a
+      = if ko `kcmp` kn
+         then let (vs, ws) = go ko as
+                 in (v : vs, ws)
+         else let (vs, ws) = go kn as
+                 in ([], (kn, v :| vs) : ws)
+
+-- | Variation of 'groupByKey' that matches mapped keys on equality.
+--
+-- >>> groupByKeyOn toLower id [("A", 1), ("a", 2), ("b", 3)]
+-- [("A",1 :| [2]),("b",3 :| [])]
+--
+-- @
+-- 'groupByKeyOn' f ≡ 'groupByKeyBy' ((==) `on` f)
+-- @
+groupByKeyOn
+  :: (Container t, Eq k')
+  => (k -> k') -> (Element t -> (k, v)) -> t -> [(k, NonEmpty v)]
+groupByKeyOn f = groupByKeyBy ((==) `on` f)
+
+-- | Operates like 'groupByFst', but uses the provided getter
+-- for the key to group by and for the value.
+--
+-- >>> groupByKey (\x -> (x `mod` 5, x)) [1, 6, 7, 2, 12, 11]
+-- [(1,1 :| [6]),(2,7 :| [2,12]),(1,11 :| [])]
+groupByKey
+  :: (Container t, Eq k)
+  => (Element t -> (k, v)) -> t -> [(k, NonEmpty v)]
+groupByKey = groupByKeyBy (==)
+
+-- | Operates similarly to 'group', grouping by the first element
+-- of the pair and returning that element in pair with each group.
+--
+-- >>> groupByFst [(1, "a"), (1, "b"), (2, "c"), (1, "d")]
+-- [(1,"a" :| ["b"]),(2,"c" :| []),(1,"d" :| [])]
+groupByFst :: Eq a => [(a, b)] -> [(a, NonEmpty b)]
+groupByFst = groupByKey id
+
+-- | Variation of 'group' that matches mapped valued on equality.
+--
+-- >>> groupOn toLower ["A", "a", "b"]
+-- ["A" :| ["a"],"b" :| []]
+--
+-- @
+-- 'groupOn' f ≡ 'groupBy' ((==) `on` f)
+-- @
+groupOn
+  :: (Container t, Eq a')
+  => (Element t -> a') -> t -> [NonEmpty (Element t)]
+groupOn f = groupBy ((==) `on` f) . toList