Transpose nested data structures

Question

I need a function like Vector [Vector a] -> Vector [Vector a], where everything stays within the same list-index, but (for each list-index) the i-th element of the j-th inner-vector becomes the j-th element of the i-th inner-vector. I suspect I'm making this a lot harder than it needs to be.

I think I specifically want to use this implementation of fixed-length vectors, but I could consider other options.

I have it working on a stand-in Pair class I wrote for simplicity, but my implementation relies heavily on pattern-matching against the structure of Pair, in a way that I think won't work for a Vec n.

The Pair class is barely relevant; it's similar to Vec nat2, but I only implemented what I thought I'd need:

data Pair a where
  Pair :: a -> a -> Pair a
  deriving (Read, Show, Eq)
instance Functor Pair where
  fmap f (Pair a1 a2) = Pair (f a1) (f a2)
instance Applicative Pair where
  pure a = Pair a a
  (Pair f1 f2) <*> (Pair a1 a2) = Pair (f1 a1) (f2 a2)
instance Foldable Pair where
  foldr f b (Pair a1 a2) = a1 `f` (a2 `f` b)

(!) :: Pair a -> Bool -> a
(Pair a1 a2) ! b = if b then a2 else a1

universe :: Pair Bool
universe = Pair False True

Working code:

transpose :: Pair [Pair a] -> Pair [Pair a]
transpose (Pair b1s b2s) = t' b1s b2s
  where t' ((Pair a11 a12) : a1s) ((Pair a21 a22) : a2s)
            = Pair ((Pair a11 a21) :) ((Pair a12 a22) :) <*> t' a1s a2s
        t' [] [] = pure []
        -- A runtime error here is ok; I have external guarantees that this won't happen.
        t' _ _ = error "Transposition requires the lists to be of equal length."

Example:

> let subject = Pair [
      Pair (1, 1, 1) (1, 1, 2),
      Pair (1, 2, 1) (1, 2, 2),
      Pair (1, 3, 1) (1, 3, 2)
  ] [
      Pair (2, 1, 1) (2, 1, 2),
      Pair (2, 2, 1) (2, 2, 2),
      Pair (2, 3, 1) (2, 3, 2)
  ]
> transpose subject
↠ Pair [
      Pair (1,1,1) (2,1,1),
      Pair (1,2,1) (2,2,1),
      Pair (1,3,1) (2,3,1)
  ] [
      Pair (1,1,2) (2,1,2),
      Pair (1,2,2) (2,2,2),
      Pair (1,3,2) (2,3,2)
  ]

With some struggle I can almost get an implementation that exclusively uses list/functor/applicative/index-able operators instead of pattern matching (and should therefore work for a Vec n), but I keep getting stuck.

How can I re-write my transpose function to work on some other structure like Pair, but with parametric size?

user_791604 · Answer

I think I'd use ZipList and a new, analogous ZipVector to achieve this. We'll start with a more standard transpose:

transposeVV :: Vector (Vector a) -> Vector (Vector a)
transposeVV = getZipVector . traverse Finite

To be able to use this, we'll need to be able to put the two Vector's next to each other; we'll do this via two more transposes that are inverses of each other:

transposeVL :: Vector [a] -> [Vector a]
transposeVL = getZipList . traverse ZipList

transposeLV :: [Vector a] -> Vector [a]
transposeLV = getZipVector . traverse Finite

The thing you want is now a conjugation of these:

transposeVlV :: Vector [Vector a] -> Vector [Vector a]
transposeVlV = transposeLV . fmap transposeVV . transposeVL

Okay, well, I hid some of the functionality away behind ZipVector. What does that look like? It's easy enough; however, one twist compared to ZipList is that we can't easily construct an infinite vector. No big deal, we'll fake it. (Your length-indexed vectors won't need to pull this trick, obviously.)

data ZipVector a
    = Inf a
    | Finite (Vector a)
    deriving Functor

getZipVector :: ZipVector a -> Vector a
getZipVector (Inf a) = V.singleton a -- hmmmm...
getZipVector (Finite as) = as

instance Applicative ZipVector where
    pure = Inf
    Inf f <*> Inf x = Inf (f x)
    Inf f <*> Finite xs = Finite (f <$> xs)
    Finite fs <*> Inf x = Finite (($x) <$> fs)
    Finite fs <*> Finite xs = Finite (V.zipWith ($) fs xs)

We can try it out now. First let's define your subject:

subject :: Vector [Vector (Int, Int, Int)]
subject =
    V.generate 2 $ \i ->
        flip map [0..2] $ \j ->
            V.generate 2 $ \k ->
                (i+1, j+1, k+1)

Then, in ghci:

> transposeVlV subject
[[[(1,1,1),(2,1,1)],[(1,2,1),(2,2,1)],[(1,3,1),(2,3,1)]],[[(1,1,2),(2,1,2)],[(1,2,2),(2,2,2)],[(1,3,2),(2,3,2)]]]

Transpose nested data structures

Working code:

Example:

1 Answers