Nice data-structure for grouping?

Learn
4

If you want a “simple Haskell” solution, you could use generic functions that allow you to specify the grouping criteria.

#!/usr/bin/env cabal
{- cabal:
build-depends:
    base ^>=4.18
  , containers ^>= 0.6
-}
{- project:
with-compiler: ghc-9.6.3
-}

module Main (main) where

import Data.Map.Lazy (Map)
import qualified Data.Map.Lazy as Map
import Data.Word (Word16, Word32, Word8)

------------------------------------------------------------------------------

groupToMapBy
  :: Ord a
  => (x -> a)
  -> [x]
  -> Map a [x]
groupToMapBy f = Map.fromListWith (++) . map (liftA2 (,) f pure)

groupToMapBy2
  :: (Ord a, Ord b)
  => (x -> a)
  -> (x -> b)
  -> [x]
  -> Map a (Map b [x])
groupToMapBy2 f g = fmap (groupToMapBy g) . groupToMapBy f

groupToMapBy3
  :: (Ord a, Ord b, Ord c)
  => (x -> a)
  -> (x -> b)
  -> (x -> c)
  -> [x]
  -> Map a (Map b (Map c [x]))
groupToMapBy3 f g h = fmap (groupToMapBy2 g h) . groupToMapBy f

------------------------------------------------------------------------------

data Atom
  = Atom
    { x :: Word8
    , y :: Word16
    , z :: Word32
    , m :: Int
    }
  deriving Show

demoAtoms :: [Atom]
demoAtoms =
    [ Atom 1 1 1 1, Atom 1 1 2 2, Atom 1 2 1 3, Atom 1 2 2 4, Atom 2 1 1 5
    , Atom 2 2 1 6, Atom 3 4 5 7, Atom 3 4 5 8
    ]

main :: IO ()
main = print $ groupToMapBy3 x y z demoAtoms
3 Likes