Haskell create vector with subvectors using indexes

Question

I'm trying to create a vector with subvectors consisting of elements taken out from another vector using a vector of sub-vector indexes. Each element in b corresponds to the sub-vector-index the elements in a should have when put into c.

import Data.Vector
let a = fromList [9,2,3,7,4,1,8,5]
let b = fromList [3,3,2,0,1,1,2,2]
let c = fromList [ a ! k | k <- b ]
Expected c = [[7],[4,1],[3,8,5],[9,2]]

I'm a bit stuck, getting the error

"Could not match expected type [Int] with actual type Vector Integer in stmt list comprehension k <- b"

Answer 1

This doesn't work since b is a Vector, not a list:

k <- b

However, this can work:

[ ... | k <- toList b ]

Next, the type of a and b is Vector Integer, and the ! operator takes an Int. So you need to convert the index using fromInteger:

let c = fromList [ a ! fromInteger k | k <- toList b]

Update

Here is a way to perform the transformation without repeated passes over the arrays:

import Data.List

fst3  (b,_,_) = b
third (_,_,a) = a

doit :: Vector Int -> Vector Int -> [[Int]]
doit av bv = [ map third g | g <- groups ]
  where
    triples = zip3 (V.toList bv) [1..] (V.toList av)
    groups = groupBy (\s t -> fst3 s == fst3 t) $ sort triples

This is basically a Schwartzian Transform with a groupBy added after the sorting step. The sorting on the triples is done in the canonical way - a lex sort on the first coordinate followed by the second coordinate followed by the third coordinate.

There are other ways to write the expression for groups:

import Data.Funcition (on)
import GHC.Exts (groupWith)

    ...
    groups = groupBy (on (==) fst3) $ sort triples
    groups = groupWith fst3 triples

Note that groupBy requires that the triples be sorted whereas groupWith doesn't.

Answer 2

With lists this seems to be the logic

> map (map snd) $ groupBy ((==) `on` fst) $ sortBy (comparing fst) $ zip b a

[[7],[4,1],[3,8,5],[9,2]]

Answer 3

With the help of ErikR I came up with this:

let c = fromList [fromList [a ! i | i <- [0..Data.Vector.length b-1], (b ! i)==j] | j <- [0..Data.Vector.maximum(b)]]

It works but it's not pretty, any better?

Answer 4

It looks like what you want would probably be

accumulate (flip (:)) (replicate n []) (zip b a)

...although you're going to have to explicitly calculate n, perhaps as maximum b + 1.