CODEWITHSUNDEEP
haskell
squirrels
squirrels

Reputation: 323

What's the most efficient way to take up to the last n elements of a list

To clarify : I want all the elements but the last n

Reading some answers on stackoverflow I get the impression that using length on lists is inadvisable. So is there better way to do this than take (length xs - n) xs?

Upvotes: 9

Views: 187

Answers (1)

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 476729

Yes. You first drop n elements from the list, then walk over the two lists concurrently and when the you hit the end of the list with the dropped version of the list, then you return the list of the "iterator" over the entire list:

takeLast :: Int -> [a] -> [a]
takeLast n ls = go (drop n ls) ls
    where go [] ls = ls
          go (_:xs) ~(_:ys) = go xs ys

This thus works because the first "iterator" is n steps ahead. So if it hits the end of the list, the second iterator is n steps behind the end of the list.

For example:

Prelude> takeLast 2 [1,4,2,5,1,3,0,2]
[0,2]
Prelude> takeLast 3 [1,4,2,5,1,3,0,2]
[3,0,2]
Prelude> takeLast 4 [1,4,2,5,1,3,0,2]
[1,3,0,2]
Prelude> takeLast 5 [1,4,2,5,1,3,0,2]
[5,1,3,0,2]

We can also drop the last n elements in a similar way:

dropLast :: Int -> [a] -> [a]
dropLast n ls = go (drop n ls) ls
    where go [] _ = []
          go (_:xs) ~(y:ys) = y : go xs ys

For example:

Prelude> dropLast 2 [1,4,2,5,1,3,0,2]
[1,4,2,5,1,3]
Prelude> dropLast 3 [1,4,2,5,1,3,0,2]
[1,4,2,5,1]
Prelude> dropLast 4 [1,4,2,5,1,3,0,2]
[1,4,2,5]
Prelude> dropLast 5 [1,4,2,5,1,3,0,2]
[1,4,2]

If we operate on an infinite list, then dropLast will still yield elements, whereas if we would use take (length ls - n) ls, it would get stuck in an infinite loop.

We can, as @DanielWagner says use zipWith for this:

dropLast :: Int -> [a] -> [a]
dropLast n xs = zipWith const xs (drop n xs)

Here we let zipWith iterate over both lists, and we use const to each time return the element of the first list.

Upvotes: 8

Related Questions