CODEWITHSUNDEEP
haskell
bennofs
bennofs

Reputation: 11973

Insert into Data.Set and check if element exists at the same time

Is there an efficient way to insert a value into a Data.Set while at the same time checking if that value was already a member of the set?

If there isn't, is there any particular reason such a function would be impossible to write with the current implementation of sets in the containers library?

Upvotes: 6

Views: 1660

Answers (2)

felix-eku
felix-eku

Reputation: 2213

It is actually possible to write such a function by slightly changing the Set.insert function. I decided to return a Maybe (Set a), so the function only creates a new Set if the element did not alredy exist. It would be equally well possible to write a function with (Bool, Set a) as return type.

insertMember :: Ord a => a -> Set a -> Maybe (Set a)
insertMember = go
  where
    go :: Ord a => a -> Set a -> Maybe (Set a)
    STRICT_1_OF_2(go)
    go x Tip = Just $ singleton x
    go x (Bin sz y l r) = case compare x y of
        LT -> do 
           l' <- go x l
           return $ balanceL y l' r
        GT -> do 
           r' <- go x r
           return $ balanceR y l 
        EQ -> Nothing

Upvotes: 2

bheklilr
bheklilr

Reputation: 54068

You can do this with O(log n) complexity by taking advantage of the fact that size is O(1), and just compare before and after:

-- | Inserts a value into the Set.  If the value was not already in the set,
-- | then True is returned, otherwise False
insertIfMissing :: Ord a => a -> Set a -> (Set a, Bool)
insertIfMissing a s = (newSet, missing)
    where
        newSet = Set.insert a s
        oldSize = Set.size s
        newSize = Set.size newSet
        missing = oldSize < newSize

And if you aren't interested in whether it was already present, then this shouldn't compute the missing part thanks to laziness.

Upvotes: 7

Related Questions