Insert value into multiple locations in a list

Is it intended that insertAt 999 [1] [1..10] inserts 999 at index 1, but insertAt 999 [0] [1..10] does nothing. Shouldn't it insert a value at index 0?

Also, it looks like certain other insertions don't work correctly:

> insertAt 999 [2,5,6,7] [1..10]
[1,2,999,3,4,999,5,999,6,7,8,9,10]

Note that this inserts only at indices 2, 5, and 7 and appears to skip index 6.

My guess is that you probably wanted something more like:

insertAt :: a -> [Int] -> [a] -> [a]
insertAt x indices xs = go 0 xs
  where
    go _ [] = []
    go j (y:ys)
      | j `elem` indices = x : go (j+1) (y:ys)
      | otherwise        = y : go (j+1) ys

which handles these cases correctly:

λ> insertAt 999 [0] [1..10]
[999,1,2,3,4,5,6,7,8,9,10]
λ> insertAt 999 [2,5,6,7] [1..10]
[1,2,999,3,4,999,999,999,5,6,7,8,9,10]

It doesn't insert anything for insertAt 999 [10] [1..10], but if you want that to insert at the end, then you need a slight modification:

insertAt :: a -> [Int] -> [a] -> [a]
insertAt x indices xs = go 0 xs
  where
    go j yall   | j `elem` indices =  x : go (j+1) yall
    go j (y:ys)                    =  y : go (j+1) ys
    go _ []                        = []

which works like so:

λ> insertAt 999 [0] [1..10]
[999,1,2,3,4,5,6,7,8,9,10]
λ> insertAt 999 [2,5,6,7] [1..10]
[1,2,999,3,4,999,999,999,5,6,7,8,9,10]
λ> insertAt 999 [10] [1..10]
[1,2,3,4,5,6,7,8,9,10,999]
λ> insertAt 999 [10,11,12] [1..10]  -- multiple consecutive indices off the end
[1,2,3,4,5,6,7,8,9,10,999,999,999]

I think this version is quite readable. The main efficiency problem is that, if indices is more than a few elements, the repeated linear traversals of the elem function through the indices list will be slow. A second efficiency problem is that if all the insertions are near the beginning, you will end up traversing the entire input list even after all the insertions have been performed.

You can improve the performance substantially by using an IntSet. (A plain Set from Data.Set would work, too, but IntSets are even faster for sets of Ints.)

import qualified Data.IntSet as IntSet
insertAt :: a -> [Int] -> [a] -> [a]
insertAt x indices = go 0
  where
    indices' = IntSet.fromList indices
    go j yall | j `IntSet.member` indices' =  x : go (j+1) yall
    go j (y:ys)                            =  y : go (j+1) ys
    go _ []                                = []

You could also pass the IntSet as a parameter to go and delete indices as they are used. This might be faster, but you'd need to benchmark it to be sure:

import qualified Data.IntSet as IntSet
insertAt :: a -> [Int] -> [a] -> [a]
insertAt x indices = go 0 (IntSet.fromList indices)
  where
    go j indices' yall | j `IntSet.member` indices' =  x : go (j+1) (IntSet.delete j indices') yall
    go j indices' (y:ys)                            =  y : go (j+1) indices' ys
    go _ _ []                                       = []

This would also allow you to stop inserting when the list of indices is empty, but again, only benchmarking will tell whether or not this is faster:

import qualified Data.IntSet as IntSet
insertAt :: a -> [Int] -> [a] -> [a]
insertAt x indices = go 0 (IntSet.fromList indices)
  where
    go _ indices' yall | IntSet.null indices'       =  yall
    go j indices' yall | j `IntSet.member` indices' =  x : go (j+1) (IntSet.delete j indices') yall
    go j indices' (y:ys)                            =  y : go (j+1) indices' ys
    go _ _ []                                       = []

Personally, I think I probably would have approached it by writing a go function that operates on the current index and a sorted list of indices:

insertAt x indices = go 0 (sort indices)

This way, we can directly compare the current index with the next insertion point in the sorted list:

go j (i:is) (y:ys) | i == j = ...

This variant would look something like:

insertAt :: a -> [Int] -> [a] -> [a]
insertAt x indices = go 0 (sort indices)
  where
    go _ []     yall              = yall
    go j (i:idxs) yall   | j == i = x : go (j+1) idxs yall
                         | j >  i = go j idxs yall -- delete duplicate indices
    go j idxs     (y:ys)          = y : go (j+1) is ys
    go _ _        []              = []

This avoids those efficiency problems I mentioned earlier without having to switch to a different data structure. Sorting the index list and consuming it as insertions are made will avoid any traversals through the index list, since you're only ever looking at the list's head. Also, when the index list is empty, the go _ [] yall = yall case will immediately return the remaining list without further processing.

I'm not sure that there's any clever, non-recursive method that uses folds or other functions from Data.List and ends up being better than the above. You can replace the recursion with unfoldr. For example, given the version from above:

insertAt :: a -> [Int] -> [a] -> [a]
insertAt x indices xs = go 0 xs
  where
    go j yall   | j `elem` indices =  x : go (j+1) yall
    go j (y:ys)                    =  y : go (j+1) ys
    go _ []                        = []

you can rewrite this as:

insertAt :: a -> [Int] -> [a] -> [a]
insertAt x indices lst = unfoldr step (0, lst)
  where step (j, lst)  | j `elem` indices = Just (x, (j+1, lst))
        step (j, y:ys)                    = Just (y, (j+1, ys))
        step (_, [])                      = Nothing

and the other versions can be rewritten similarly. This doesn't particularly improve readability, and I doubt it's more efficient.

Another completely different approach might be to intercalate the insertion value into a sequence of sublists generated from the indexes. So, for example, given the indexes [2,5,6,7], you might want to first break the input list [1..10] up into sublists between the insertions:

[[1,2],[3,4],[],[],[5,6,7,8,9,10]]

It might be easiest to first convert a sorted list of indices into a list of sizes for all the sublists except the last:

indicesToChunks :: [Int] -> [Int]
indicesToChunks indices = zipWith (\i j -> i-j-1) indices (-1 : indices)

giving:

> indicesToChunks [2,5,6,7]
[2,2,0,0]

and then write a function to break a list into sublists based on these sublist sizes. It's pretty easy to write recursively:

chunks :: [Int] -> [a] -> [[a]]
chunks [] lst = [lst]
chunks (n:ns) lst = a : chunks ns b
  where (a,b) = splitAt n lst

Then, you can write:

insertAt :: a -> [Int] -> [a] -> [a]
insertAt x indices = intercalate [x] . chunks (indicesToChunks (sortu indices))
  where -- unique elements in sorted order
        sortu = map head . group . sort

The definition of sortu is weird, but it's a reasonable method of sorting and removing duplicates. (The nub function in Data.List removes duplicates, but doesn't do it efficiently for sorted lists, the way map head . group does.)

Note that this version behaves a little differently than the others, if there are insertions at the end of the list that skip indexes. The earlier versions all insert only one 999 at the end of the list in this example:

> insertAt 999 [10,1000] [1..10]
[1,2,3,4,5,6,7,8,9,10,999]

treating the 1000 as an out-of-bounds index. However, the intercalate-based version inserts two 999s:

> insertAt 999 [10,1000] [1..10]
[1,2,3,4,5,6,7,8,9,10,999,999]