Yet another implementation (might be the same?)

Partition2.hs

@@ -0,0 +1,62 @@
+module Partition2 where
+
+import Data.Foldable as F
+import Data.List as L
+import Data.Set as S
+import Data.Tuple (swap)
+import Data.Map.Strict as M
+import Data.Vector as V
+
+type Class = Int
+data Partition a = Partition (Map a Class) (Vector (Set a))
+  deriving Show
+
+size :: Partition a -> Int
+size (Partition _ v) = V.length v
+
+elems :: Ord a => Partition a -> Set a
+elems (Partition _ v) = V.foldr S.union S.empty v
+
+flat :: Vector (Set a) -> [(Int, a)]
+flat v = F.concatMap (\(i, s) -> foldMap (\x -> [(i, x)]) s) $ indexed v
+
+-- Ord on o is used to determine classes, Ord on a is used to
+-- store the elements
+partitionWith :: (Ord o, Ord a) => (a -> o) -> [a] -> Partition a
+partitionWith f ls = Partition map groups
+  where
+    map = M.fromList . fmap swap $ flat groups
+    groups = V.fromList . fmap S.fromList . groupOn f . sortOn f $ ls
+
+groupOn :: Eq b => (a -> b) -> [a] -> [[a]]
+groupOn f ls = L.groupBy (\a b -> f a == f b) ls
+
+instance Ord a => Eq (Partition a) where
+  p1@(Partition m1 v1) == p2@(Partition m2 v2)
+    | Partition.size p1 /= Partition.size p2 = False
+    | Partition.elems p1 /= Partition.elems p2 = False
+    | otherwise = wellDefined v1 m2 && wellDefined v2 m1
+    where wellDefined v m = V.all (\x -> S.size x == 1) $ imap (\i s -> S.map (\x -> M.lookup x m) s) v
+
+fromJustToSet Nothing = S.empty
+fromJustToSet (Just s) = S.singleton s
+
+dseflatten :: Ord a => Set (Set a) -> Set a
+dseflatten = S.unions . S.toList
+
+dseconcatMap :: (Ord a, Ord b) => (a -> Set b) -> Set a -> Set b
+dseconcatMap f s = dseflatten (S.map f s)
+
+
+--{-# LANGUAGE ViewPatterns, TupleSections #-}
+--import Data.Foldable
+--
+--M.lookup
+--
+--let p = partitionWith (== 1) [1,2,3,4]
+--let p2 = partitionWith (/= 1) [1,2,3,4]
+--let p3 = partitionWith (== 2) [1,2,3,4]
+--print p
+--print p2
+--p == p2
+--p == p3