Apply a function to the nth element of a list

1
votes

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.

4
haskell
Take a look at lenses. – jamshidh
Taking a look at lenses may not be the best idea for somebody new to Haskell. I think that what you have here is fine. Are you looking for something more efficient? More idiomatic? A one-liner? – Alexander Vieth
@jamshidh I had a quick look at lenses, but I would prefer to be able to do this without any packages yet. – b3036667
@AlexanderVieth Are those options mutually exclusive? It's just, in my limited experience with Haskell, the approach a beginner takes is usually the 'wrong' approach for Haskell. – b3036667
No they are not mutually exclusive, just wondering what your goal is. doSomething is certainly correct, and I could understand it very quickly, so I see nothing 'wrong' with it except for performance considerations (++ is potentially slow). – Alexander Vieth

4 Answers

4
votes

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

> import Control.Lens.Combinators
> 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
3
votes

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]
2
votes

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]
2
votes

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.