diff --git a/src/Data/Map.purs b/src/Data/Map.purs index abbbb179..376abfaf 100644 --- a/src/Data/Map.purs +++ b/src/Data/Map.purs @@ -32,12 +32,14 @@ module Data.Map , unions , size , mapWithKey + , split ) where import Prelude import Data.Eq (class Eq1) import Data.Foldable (foldl, foldMap, foldr, class Foldable) import Data.List (List(..), (:), length, nub) +import Data.List.Lazy as LL import Data.Maybe (Maybe(..), maybe, isJust, fromMaybe) import Data.Monoid (class Monoid) import Data.Ord (class Ord1) @@ -451,3 +453,11 @@ mapWithKey :: forall k v v'. (k -> v -> v') -> Map k v -> Map k v' mapWithKey _ Leaf = Leaf mapWithKey f (Two left k v right) = Two (mapWithKey f left) k (f k v) (mapWithKey f right) mapWithKey f (Three left k1 v1 mid k2 v2 right) = Three (mapWithKey f left) k1 (f k1 v1) (mapWithKey f mid) k2 (f k2 v2) (mapWithKey f right) + +-- | Divide into two maps of keys less and greater than the provided argument. +split :: forall k v. Ord k => k -> Map k v -> { less :: Map k v, greater :: Map k v } +split k = mapify <<< LL.span (\(Tuple k' _) -> k' < k) <<< toAscUnfoldable + where + mapify {init: ls, rest: gs} = + {less: fromFoldable ls, + greater: fromFoldable $ LL.dropWhile (\(Tuple k' _) -> k' == k) gs} diff --git a/test/Test/Data/Map.purs b/test/Test/Data/Map.purs index 26612d8d..6c42d5fd 100644 --- a/test/Test/Data/Map.purs +++ b/test/Test/Data/Map.purs @@ -274,3 +274,9 @@ mapTests = do toList = M.toUnfoldable :: forall k v. M.Map k v -> List (Tuple k v) resultViaLists = m # toList # map (\(Tuple k v) → Tuple k (f k v)) # M.fromFoldable in resultViaMapWithKey === resultViaLists + + log "split is correct" + quickCheck $ \(TestMap m :: TestMap String Int) s -> let + {less: ls, greater: gs} = M.split s m + in (all (_ < s) $ M.keys ls) && (all (_ > s) $ M.keys gs) && + M.size ls + M.size gs == M.size (M.delete s m)