Skip to content

Commit eb4be61

Browse files
committed
Add size aware intersection and difference functions
This commit adds versions of 'intersection, intersectionWith, intersectionWithKey, difference, differenceWith' that track size to be used later in this PR.
1 parent 28d25e4 commit eb4be61

File tree

1 file changed

+80
-0
lines changed

1 file changed

+80
-0
lines changed

Data/HashMap/Base.hs

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1213,6 +1213,21 @@ difference a b = foldlWithKey' go empty a
12131213
_ -> m
12141214
{-# INLINABLE difference #-}
12151215

1216+
-- | /O(n*log m)/ Difference of two maps. Return elements of the first map
1217+
-- not existing in the second. Returns a tuple with the size of the resulting
1218+
-- hashmap, and the hashmap itself.
1219+
differenceInternal
1220+
:: (Eq k, Hashable k)
1221+
=> HashMap k v
1222+
-> HashMap k w
1223+
-> (Int, HashMap k v)
1224+
differenceInternal a b = foldlWithKey' go (0, empty) a
1225+
where
1226+
go p@(sz, m) k v = case lookup k b of
1227+
Nothing -> let (s, m') = insertInternal k v m in (sz + s, m')
1228+
_ -> p
1229+
{-# INLINABLE differenceInternal #-}
1230+
12161231
-- | /O(n*log m)/ Difference with a combining function. When two equal keys are
12171232
-- encountered, the combining function is applied to the values of these keys.
12181233
-- If it returns 'Nothing', the element is discarded (proper set difference). If
@@ -1225,6 +1240,27 @@ differenceWith f a b = foldlWithKey' go empty a
12251240
Just w -> maybe m (\y -> insert k y m) (f v w)
12261241
{-# INLINABLE differenceWith #-}
12271242

1243+
-- | /O(n*log m)/ Difference with a combining function. When two equal keys are
1244+
-- encountered, the combining function is applied to the values of these keys.
1245+
-- If it returns 'Nothing', the element is discarded (proper set difference). If
1246+
-- it returns (@'Just' y@), the element is updated with a new value @y@.
1247+
-- Returns a tuple with the size of the resulting hashmap, and the hashmap
1248+
-- itself.
1249+
differenceWithInternal
1250+
:: (Eq k, Hashable k)
1251+
=> (v -> w -> Maybe v)
1252+
-> HashMap k v
1253+
-> HashMap k w
1254+
-> (Int, HashMap k v)
1255+
differenceWithInternal f a b = foldlWithKey' go (0, empty) a
1256+
where
1257+
go p@(sz, m) k v = case lookup k b of
1258+
Nothing -> let (s, m') = insertInternal k v m in (sz + s, m')
1259+
Just w -> maybe p (\y -> let (s, m') = insertInternal k y m
1260+
in (sz + s, m'))
1261+
(f v w)
1262+
{-# INLINABLE differenceWithInternal #-}
1263+
12281264
-- | /O(n*log m)/ Intersection of two maps. Return elements of the first
12291265
-- map for keys existing in the second.
12301266
intersection :: (Eq k, Hashable k) => HashMap k v -> HashMap k w -> HashMap k v
@@ -1235,6 +1271,22 @@ intersection a b = foldlWithKey' go empty a
12351271
_ -> m
12361272
{-# INLINABLE intersection #-}
12371273

1274+
-- | /O(n*log m)/ Intersection of two maps. Return elements of the first
1275+
-- map for keys existing in the second.
1276+
--- The result is a tuple with the size of the resulting hashmap, and the
1277+
-- hashmap itself.
1278+
insersectionInternal
1279+
:: (Eq k, Hashable k)
1280+
=> HashMap k v
1281+
-> HashMap k w
1282+
-> (Int, HashMap k v)
1283+
insersectionInternal a b = foldlWithKey' go (0, empty) a
1284+
where
1285+
go p@(sz, m) k v = case lookup k b of
1286+
Just _ -> let (s, m') = insertInternal k v m in (sz + s, m')
1287+
_ -> p
1288+
{-# INLINABLE insersectionInternal #-}
1289+
12381290
-- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps
12391291
-- the provided function is used to combine the values from the two
12401292
-- maps.
@@ -1247,6 +1299,20 @@ intersectionWith f a b = foldlWithKey' go empty a
12471299
_ -> m
12481300
{-# INLINABLE intersectionWith #-}
12491301

1302+
-- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps
1303+
-- the provided function is used to combine the values from the two
1304+
-- maps.
1305+
-- The result is a tuple with the size of the resulting hashmap, and the
1306+
-- hashmap itself.
1307+
insersectionWithInternal :: (Eq k, Hashable k) => (v1 -> v2 -> v3) -> HashMap k v1
1308+
-> HashMap k v2 -> (Int, HashMap k v3)
1309+
insersectionWithInternal f a b = foldlWithKey' go (0, empty) a
1310+
where
1311+
go p@(sz, m) k v = case lookup k b of
1312+
Just w -> let (s, m') = insertInternal k (f v w) m in (sz + s, m')
1313+
_ -> p
1314+
{-# INLINABLE insersectionWithInternal #-}
1315+
12501316
-- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps
12511317
-- the provided function is used to combine the values from the two
12521318
-- maps.
@@ -1259,6 +1325,20 @@ intersectionWithKey f a b = foldlWithKey' go empty a
12591325
_ -> m
12601326
{-# INLINABLE intersectionWithKey #-}
12611327

1328+
-- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps
1329+
-- the provided function is used to combine the values from the two
1330+
-- maps.
1331+
-- The result is a tuple with the size of the resulting hashmap, and the
1332+
-- hashmap itself.
1333+
intersectionWithKeyInternal :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3)
1334+
-> HashMap k v1 -> HashMap k v2 -> (Int, HashMap k v3)
1335+
intersectionWithKeyInternal f a b = foldlWithKey' go (0, empty) a
1336+
where
1337+
go p@(sz, m) k v = case lookup k b of
1338+
Just w -> let (s, m') = insertInternal k (f k v w) m in (sz+s, m')
1339+
_ -> p
1340+
{-# INLINABLE intersectionWithKeyInternal #-}
1341+
12621342
------------------------------------------------------------------------
12631343
-- * Folds
12641344

0 commit comments

Comments
 (0)