How to pick multiple elements from a list in Haskell

Viewed 145

In Haskell I know how I can pick one element from a list by index :

head (drop idx myList)

But I have a list of indices myIndices (like [1,3,5]) and want to return a list of the elements in myList at the index positions in myIndices.

In a non-functional language I would build up the list over a loop, but how do I "build" the new list in Haskell? Can I use a filter?

2 Answers

First let me add my standard remark that if you find yourself thinking about indices in a list you should reconsider your entire approach. It's better to think of lists as streams of values, where you never really care about absolute position but only about the values and possibly their relative order. If you really need indices, you should probably be using Vectors instead of lists.

But ok. If you would use a loop for some task in an imperative language, the obvious approach is always to use recursion in a functional language. Often it's not the most concise, reliable or efficient option, but it's what is always possible and you should certainly practice how to do it.

Basically you write a go function that has an i argument which is “incremented” for each recursive call much like the counter variable in a for loop, and the result is just a list which is build by purely prepending elements when i is in the indices list. Do implement this as an exercise.

But, yes, you could also use filter instead. The trick is that instead of manually maintaining a counter variable, you just add that information to the list as it were. In fact that's something often done even in Python

for i,x in enumerate(myList):
    ...

In Haskell, thanks to lazy evaluation, enumerate doesn't need any special iterable magic but can be expressed as a simple list zip: in fact the above Python loop could be written in e.g. the IO monad as

forM_ (zip [0..] myList) $ \(i,x) -> do
    ...

But here you don't really want to loop over the list, only, indeed, filter on it. Namely, you want to retain all the elements whose index is in the list of requested indices:

  [ x | (i,x) <- zip [0..] myList, i`elem`myIndices ]

Note that elem has bad performance though: it'll go through all the requested indices for all elements in myList, whereas if myIndices is sorted it would be enough to check only the head element. That's what you could do better in a recursive solution, or alternatively by using the State monad to pop the used indices.

But a rather nicer approach is to create a “masking” list from the list of indices, i.e. turning [1,3,5] into [False, True, False, True, False, True]. This could again be done with manual recursion, or with a trick like this:

indicesAsMask :: [Int]  -- Must be strictly sorted
              -> [Bool]
indicesAsMask ixs@(0:_) = True : do
   (il, ir) <- zip ixs $ tail ixs
   replicate (ir - il - 1) False ++ [True]
indicesAsMask ixs = False : tail (indicesAsMask $ 0:ixs)

Then you can simply zip your list together with the mask, and filter out those elements whose mask value is True.

If you know the indices are not out of bound, you can go with

fmap (myList !!) myIndices

This requires that myList be indexable at every index in myIndices. If not, you get a runtime exception:

myList = ["a", "b", "c", "d"]
myIndices = [1,2,3,10]
fmap (myList !!) myIndices
["b","c", "d","*** Exception: Prelude.!!: index too large

If you want to be robust in that respect, you can take a slightly different approach:

concat $ fmap (\i -> take 1 $ drop i $ myList) myIndices

Here each index i is turned into a list containing the one corresponding element from myList or an empty list if myList is not long enough; the resulting list of lists is flattened by concat.

The above can be written more compactly using the monadic bind operator >>=:

myIndices >>= (\i -> take 1 $ drop i $ myList)

The drawback of this approach is that it runs take 1 $ drop i $ myList for every single index, meaning that you traverse myList multiple times, thus giving a complexity O(m*n) for m indices and n elemnts in the list, whereas you could do it only once, in principle, as shown in the other answer. (More details here.) Whether such a complexity is acceptable, clearly depends how big m*n is.

Related