Group list of words by occurrence / inverse index of strings I'm currently working through Coding the matrix, though I'm trying to work through it using Haskell rather than Python. One of the exercises has you write a function where you take a list of strings, and then transform them into an inverse index of the words contained in each string, along with the index of the originating string:

&gt; makeInverseIndex ["hello world", "world test", "hello test"]
fromList [("hello",fromList [0,2]),("test",fromList [1,2]),("world",fromList [0,1])]


You then implement an orSearch and an andSearch, which returns the indexes of the strings that match the search:

&gt; orSearch ["hello", "world"] $ makeInverseIndex ["hello world", "world test", "hello test"]
fromList [0,1,2]

&gt; andSearch ["hello", "world"] $ makeInverseIndex ["hello world", "world test", "hello test"]
fromList [0]


Here is my attempt in Haskell:

import Data.List as L
import Data.Map.Strict as M
import Data.Set as S

makeInverseIndex :: [String] -&gt; Map String (Set Int)
makeInverseIndex = L.foldr (unionWith S.union . uncurry group) M.empty . zipWithIndex . L.map words
  where
    zipWithIndex :: [a] -&gt; [(Int, a)]
    zipWithIndex xs = zip [0..length xs] xs

    group :: Ord k =&gt; v -&gt; [k] -&gt; Map k (Set v)
    group n = M.fromList . L.map (\w -&gt; (w, S.singleton n))

orSearch :: [String] -&gt; Map String (Set Int) -&gt; Set Int
orSearch words =
  M.foldr S.union S.empty . pick words

andSearch :: [String] -&gt; Map String (Set Int) -&gt; Set Int
andSearch words index =
  M.foldr S.intersection (S.fromList [0..length index-1]) . pick words $ index

pick :: Ord k =&gt; [k] -&gt; Map k v -&gt; Map k v
pick keys m = 
  restrictKeys m $ S.fromList keys
Prefer qualified imports for containers

The containers modules contain several functions that have the same names as their list counterpart. Therefore, they are usually included as qualified modules:

import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import           Data.Set (Set)
import qualified Data.Set as S


The types are imported unqualified to make the type signatures easier to read.

Next, it's a easier to exchange your types later if you provide a type synonym:

type MultiMap k v = Map k (Set v)
type WordMap      = MultiMap String Int


Now, let's have a look at your functions. zipWithIndex isn't optimal because it traverses xs twice. However the result of zip is as long as the shorter of the two lists. We can therefore simply write

zipWithIndex xs = zip [0..] xs


Note that you don't need a type signature on those local functions. Indeed, they can be misleading, because the a in zipWithIndex is not related to the a in the outer function.

Don't shadow library function names

group is a name that's already imported via Data.List. Since we now import our containers as qualified, we can simply provide our own singleton function to get rid of group:

singleton :: k -&gt; v -&gt; MultiMap k v
singleton k v = M.singleton k (S.singleton v)


Our makeInverseFunction would now look like this:

singleton :: k -&gt; v -&gt; MultiMap k v
singleton k v = M.singleton k (S.singleton v)

makeInverseIndex :: [String] -&gt; WordMap
makeInverseIndex = foldr (M.unionWith S.union . uncurry insert) M.empty . zipWithIndex . L.map words
  where
    zipWithIndex = zip [0..]

    insert v = foldMap (flip singleton v)


Prefer functions that provide your functionality already

However, there's a function to convert a list of maps into a single map, , unionsWith:

makeInverseIndex :: [String] -&gt; WordMap
makeInverseIndex = M.unionsWith S.union . concat . zipWith go [0..] . map words
  where
    go index ws = map (flip singleton index) ws


I admit that go is a bad name in that context. indexer might be a better one. By the way, we cannot use foldMap or mconcat here, since that wouldn't merge the map values.

Try to relax your type signatures*

orSearch is fine, although you could relax its type. Also, words is against a function name. You can use ws safely in this context.

orSearch :: (Ord k, Ord v) =&gt; [k] -&gt; MultiMap k v -&gt; Set v
orSearch ws =
  M.foldr S.union S.empty . pick ws


* unless it leads to performance problems or ambiguities

Prefer foldr1 instead of complicated start values

Now, andSearch is a little bit tricky. You use S.fromList [0..length index-1] in order to have a proper "zero" case. However, if the map is empty, the correct answer should be the empty set, not the complete, right?

So let's handle that case first with M.null index and then use foldr1 from Map's Foldable instance:

andSearch :: (Ord k, Ord v) =&gt; [k] -&gt; MultiMap k v -&gt; Set v
andSearch words index =
  | M.null sets = S.empty
  | otherwise   = foldr1 S.intersection sets
 where
   sets = pick words index


Other than that, well done. Keep in mind that there's a IntSet in Data.IntSet that might be more suitable for your use case.