From e5d80f3ed56bdf8685f39be246191b51cb50f3f8 Mon Sep 17 00:00:00 2001
From: John Watson <john.watson@gmx.co.uk>
Date: Thu, 28 Jul 2022 18:05:45 +0100
Subject: [PATCH] Add transpose to Data.Array (#226)

* Add transpose to Data.Array

Implementation courtesy of Jordan Martinez.

* Update CHANGELOG

* Improve layout of doc comments
---
 CHANGELOG.md              |  1 +
 src/Data/Array.purs       | 41 +++++++++++++++++++++++++++++++++++++++
 test/Test/Data/Array.purs | 14 +++++++++++++
 3 files changed, 56 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 387f4a3c..9c96efc1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ Notable changes to this project are documented in this file. The format is based
 Breaking changes:
 
 New features:
+- Added `transpose` to `Array` (#225 by @newlandsvalley and @JordanMartinez)
 
 Bugfixes:
 
diff --git a/src/Data/Array.purs b/src/Data/Array.purs
index 6bc035ab..866cc309 100644
--- a/src/Data/Array.purs
+++ b/src/Data/Array.purs
@@ -84,6 +84,7 @@ module Data.Array
   , foldMap
   , fold
   , intercalate
+  , transpose
   , scanl
   , scanr
 
@@ -781,6 +782,46 @@ fold = F.fold
 intercalate :: forall a. Monoid a => a -> Array a -> a
 intercalate = F.intercalate
 
+-- | The 'transpose' function transposes the rows and columns of its argument.
+-- | For example,
+-- |
+-- | ```purescript
+-- | transpose 
+-- |   [ [1, 2, 3]
+-- |   , [4, 5, 6]
+-- |   ] == 
+-- |   [ [1, 4]
+-- |   , [2, 5]
+-- |   , [3, 6]
+-- |   ]
+-- | ```
+-- |
+-- | If some of the rows are shorter than the following rows, their elements are skipped:
+-- |
+-- | ```purescript
+-- | transpose 
+-- |   [ [10, 11]
+-- |   , [20]
+-- |   , [30, 31, 32]
+-- |   ] == 
+-- |   [ [10, 20, 30]
+-- |   , [11, 31]
+-- |   , [32]
+-- |   ]
+-- | ```
+transpose :: forall a. Array (Array a) -> Array (Array a)
+transpose xs = go 0 []
+  where
+  go :: Int -> Array (Array a) -> Array (Array a)
+  go idx allArrays = case buildNext idx of
+    Nothing -> allArrays
+    Just next -> go (idx + 1) (snoc allArrays next)  
+   
+  buildNext :: Int -> Maybe (Array a)
+  buildNext idx = do
+    xs # flip foldl Nothing \acc nextArr -> do
+      maybe acc (\el -> Just $ maybe [el] (flip snoc el) acc) $ index nextArr idx
+
 -- | Fold a data structure from the left, keeping all intermediate results
 -- | instead of only the final result. Note that the initial value does not
 -- | appear in the result (unlike Haskell's `Prelude.scanl`).
diff --git a/test/Test/Data/Array.purs b/test/Test/Data/Array.purs
index 9325c0f8..e2d7b0fe 100644
--- a/test/Test/Data/Array.purs
+++ b/test/Test/Data/Array.purs
@@ -265,6 +265,20 @@ testArray = do
   assert $ A.modifyAtIndices [0, 2, 8] not [true,  true, true,  true] ==
                                            [false, true, false, true]
 
+  log "transpose swaps rows and columns for a regular two-dimension array"
+  assert $ A.transpose [[1,2,3], [4,5,6], [7,8,9]] ==
+                       [[1,4,7], [2,5,8], [3,6,9]] 
+  
+  log "transpose skips elements when rows don't match"
+  assert $ A.transpose [[10,11], [20], [30,31,32]] ==
+                       [[10,20,30], [11,31], [32]]
+
+  log "transpose [] == []"
+  assert $ A.transpose [] == ([] :: Array (Array Int))
+
+  log "transpose (singleton []) == []"
+  assert $ A.transpose (A.singleton []) == ([] :: Array (Array Int))                                          
+
   log "scanl should return an array that stores the accumulated value at each step"
   assert $ A.scanl (+)  0 [1,2,3] == [1, 3, 6]
   assert $ A.scanl (-) 10 [1,2,3] == [9, 7, 4]