You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
62 lines
1.8 KiB
62 lines
1.8 KiB
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
|
|
|