Bucket sort in Haskell

As a first step in learning Haskell I am solving this problem, which involves finding the Rith-ranked numbers, given some input list and many Ri. In an imperative language I would make a zeroed array c of length 200, increment c[h] for each height h, compute cumulative sums of c and then binary search c to determine the height corresponding to each given index. Because max_height is fixed, this has runtime linear in the size of input and bounded memory(*) excluding the input.

Here's my Haskell code:

max_height = 200
count_eq e = foldl (\c f -> if e == f then c + 1 else c) 0
counts heights = map (flip count_eq heights) [0..max_height]
first_gt e l = f l 0 where f (x:xs) i = if x > e then i else f xs (i+1)
solve heights indices = let accum = scanl1 (+) (counts heights) in
map (flip first_gt accum) (map (subtract 1) indices)

It is correct but slow. I would like to know how to (A) reason about and (B) improve the performance. Also (C) can I achieve the same asymptotic performance as the imperative code?

(*) assuming each c[i] fits in a machine int. I believe the runtime statement holds regardless.

• Welcome to Code Review! Good job on your first question. – SirPython Apr 7 '16 at 23:52

Maps as arrays

In counts heights you are taversing heights 200 times. It is possible to emulate array with Data.IntMap and do this in one pass:

import Data.List (foldl')
import qualified Data.IntMap.Strict as Map

count = foldl' (\m h -> Map.insertWith (+) h 1 m)
(Map.fromList [(i,0) | i <- [1..200]])

Note that I'm using strict version of left fold (see here about foldl vs foldl') and strict Map. This means that map of height counts constructed in single pass without thunks.

solve heights = map (\ix -> 1 + length (takeWhile (<ix) counts))
where
counts = scanl1 (+) $Map.elems$ count heights

I'm using 1 + length (takeWhile (<ix) counts) instead of bare recursion in first_gt.

A bit more code is required for binary search:

solve heights = map (maybe 0 snd . (Map.lookupLT countsMap))
where
counts = scanl1 (+) $Map.elems$ count heights