Haskell: chunk a list of key-value pairs on N-sized chunks, in fair order.

OrderedFairChunks.hs

module OrderedFairChunks  where 

import Data.List
import Data.Maybe
import Data.Function(on)

-- Useful type aliases
type Cons a = (a, [a]) 
type Slice k v = (k, [v]) 

--- Phase 1: mark every pair with its position/ 
zipWithIndex :: [(a,b)] -> [(a,b,Int)]
zipWithIndex = zipWith (flip app3) [0..]

--- Phase 2: split the list into the list of slices, by their keys
sliceByKeys :: Eq a => [(a,b,c)] -> [Slice a (b,c)]
sliceByKeys = unfoldr (onCons sliceOne sliceRest)

sliceOne :: Eq a => Cons (a,b,c) -> Slice a (b,c) 
sliceOne (head, tail) = (hk, dfst3 head : rest) where
  hk = fst3 head
  rest = (map dfst3) . filter ( (hk ==) . fst3 ) $ tail

sliceRest :: Eq a => Cons (a,b,c) -> [(a,b,c)]
sliceRest (head, tail) = filter ( not . (hk ==) . fst3) tail where
  hk = fst3 head

--- Phase 3: cut each group into batches of size n
chopSlices :: Int -> [Slice a (b, c)] -> [Slice a [(b, c)] ]
chopSlices n = map (mapSnd (chopSlice n))

chopSlice :: Int -> [t] -> [[t]]
chopSlice n = unfoldr (onCons chopOne chopRest) where 
  chopOne  (p,ps) = p : take (n-1) ps
  chopRest (_,ps) =     drop (n-1) ps


-- Phase 4: merge the existing groups
merge :: Ord c => [Slice a [(b, c)] ] -> [Slice a b]
merge = unfoldr (onCons merge1 merge2)

merge1 :: Cons (Slice a [(b,c)] ) -> Slice a b
merge1 (headChop,_) = mapSnd (map fst . head) headChop

merge2 :: Ord c => Cons (Slice a [(b,c)]) -> [Slice a [(b, c)]]
merge2 (head_slice, tail_slices) =
  if null chops then tail_slices else insertBy leqSlice cuthead tail_slices where
    (key, _: chops) = head_slice
    cuthead = (key, chops)

leqSlice :: Ord c => Slice a [(b, c)] -> Slice a [(b,c)] -> Ordering
leqSlice = compare  `on` firstChopNumber where 

firstChopNumber :: Slice a [(b, c)] -> c
firstChopNumber = snd . head . head . snd

---


onCons :: (Cons a -> b) -> (Cons a -> [a]) -> [a] -> Maybe (b,[a])
onCons _ _ [] = Nothing
onCons f g (x:xs) = Just (f (x,xs), g (x,xs) )

toCons :: [a] -> Maybe (a, [a])
toCons [] = Nothing
toCons (x:xs) = Just (x,xs)

-- Tuple functions: 
fst3 :: (a,b,c) -> a
fst3 (x,_,_) = x

dfst3 :: (a,b,c) -> (b,c)
dfst3 (_,y,z) = (y,z)

app3 :: (a,b) -> c -> (a,b,c) 
app3 (x,y) i = (x,y,i) 

mapSnd :: (b -> c) -> (a,b) -> (a,c)
mapSnd f (x,y) = (x, f y)

-- combine all elements
fullFunction :: Eq a => Int -> [(a, b)] -> [Slice a b]
fullFunction n = merge . chopSlices n .  sliceByKeys . zipWithIndex

choppedFunction :: Eq a => Int -> [(a, b)] -> [Slice a [(b,Int)]]
choppedFunction n = chopSlices n .  sliceByKeys . zipWithIndex

testInput :: [(Char, Int)]
testInput = [('b',4), ('a',5), ('c',3), ('c',5), ('c',4), 
         ('a',2), ('b',2), ('a',4), ('c',1), ('a',1), 
         ('a',3), ('b',1), ('b',5), ('b',3), ('c',2) ]

testExpected :: [(Char, [Int])] 
testExpected = [ ('b',[4,2,1]), ('a',[5,2,4]), 
        ('c',[3,5,4]), ('c',[1,2]), 
        ('a',[1,3]), ('b',[5,3]) ]

cyclic :: [(Char, Int) ]
cyclic = zip (cycle ['a' .. 'z']) (cycle [1..10])

main :: IO ()
main = putStrLn  ""