将函数应用于列表的第n个元素

Question

I need to be able to apply a function to the nth element of a list. For example:

> doSomething (+5) 2 [1,2,3,4,5]

should return [1,7,3,4,5]

I have a function that can do this:

doSomething :: (a -> a) -> Int -> [a] -> [a]
doSomething f n xs = ys ++ [f x] ++ zs
    where (ys, x:zs) = splitAt (n - 1) xs

but I'm new to Haskell, and so I'm sure (as with many simple functions in Haskell) there is a much better way of doing this.

Answer 1

As jamshidh indicates the lens package makes it simple to achieve this kind of task.

> over (element 2) (+5) [1..5]
[1,2,8,4,5]

This kind of operation works over any traversable, for example trees:

> import Data.Tree
> let tree = Node 1 [Node 2 [], Node 3 []]
> putStr . drawTree . fmap show $ tree
1
|
+- 2
|
`- 3
> putStr . drawTree . fmap show $ over (element 2) (+5) tree
1
|
+- 2
|
`- 8

Answer 2

If you need random access to elements of a sequence, you may not want to use a list at all. You could, for example, use Data.Vector instead:

import Data.Vector (Vector)
import qualified Data.Vector as V

modifyNth :: Int -> (a -> a) -> Vector a -> Vector a
modifyNth n f = V.imap f'
    where f' i a | i == n    = f a
                 | otherwise = a

Example use:

>>> modifyNth 2 (+5) (V.fromList [1,2,3,4,5])
fromList [1,2,8,4,5]

Answer 3

If you don't want to dive into lenses, and prefer a simple solution, you can just use list comprehensions; it runs on linear time, your list concatenations would degrade performance on large lists:

Prelude> [if i == 2 then v + 5 else v | (i, v) <- zip [1..] l]
[1,7,3,4,5]

So, doSomething would be:

Prelude> let doSomething f i l = [if p == i then f v else v | (p, v) <- zip [1..] l]
Prelude> doSomething (+5) 2 [1,2,3,4,5]
[1,7,3,4,5]

Answer 4

You can do this with some manual recursion pretty easily, and it will perform better than the splitAt version, as well as allocating fewer temporary objects than the list comprehension.

doSomething :: (a -> a) -> Int -> [a] -> [a]
doSomething _f _ [] = []
doSomething f 0 (x:xs) = f x : xs
doSomething f n (x:xs) = x : doSomething f (n - 1) xs

The cases are all pretty obvious: if the list is empty, you can't do anything, so return it. If n is 0, then just call f on it and add that to the rest of the list. Otherwise, you can put the current x at the front, and recurse with a smaller n.