Dyadic Transpose

Question

As with most APL symbols, ⍉ has different meanings when called with one argument (transpose) versus two arguments (dyadic transpose / reorder dimensions). This challenge concerns the latter, which acts similarly to numpy.moveaxis in Python or permute in MATLAB, but is more powerful.

order ⍉ A when order has distinct entries

When all members of order are distinct, order ⍉ A is equivalent to:

• numpy.moveaxis(A, tuple(range(len(A.shape)), order) in Python, or
• permute(A,order) in MATLAB. Quoting from documentation of the latter:

B = permute(A,order) rearranges the dimensions of A so that they are in the order specified by the vector order. The resulting array B has the same values as A but the order of the subscripts needed to access any particular element is rearranged as specified by order.

For example, suppose A is a 3D array, and let B ← (2 0 1)⍉A. Then B is such that B[x0,x1,x2] = A[x2,x0,x1] for all x2,x0,x1

order ⍉ A when order has repeated entries

When order has repeated entries, we take a diagonal slice of the array. For example, let A be a 2x3x4 array. B ← (0 0 1)⍉A takes a diagonal slice along A to create B such that B[x0,x1] = A[x0,x0,x1]. Note that B is a 2x4 array: if it were 3x4, we would need to set B[2, x1] = A[2, 2, x1] which would be out of bounds of A. In general the kth dimension of B will be the minimum of all A.shape[i] such that order[i] = k.

Example

Consider the dyadic transpose order⍉A where order = [2, 1, 0] and A is 3x4x5

    A =
[[[ 0  1  2  3  4]
  [ 5  6  7  8  9]
  [10 11 12 13 14]
  [15 16 17 18 19]]

 [[20 21 22 23 24]
  [25 26 27 28 29]
  [30 31 32 33 34]
  [35 36 37 38 39]]

 [[40 41 42 43 44]
  [45 46 47 48 49]
  [50 51 52 53 54]
  [55 56 57 58 59]]]

The result is the 5x4x3 array B =

[[[ 0 20 40]
  [ 5 25 45]
  [10 30 50]
  [15 35 55]]

 [[ 1 21 41]
  [ 6 26 46]
  [11 31 51]
  [16 36 56]]

 [[ 2 22 42]
  [ 7 27 47]
  [12 32 52]
  [17 37 57]]

 [[ 3 23 43]
  [ 8 28 48]
  [13 33 53]
  [18 38 58]]

 [[ 4 24 44]
  [ 9 29 49]
  [14 34 54]
  [19 39 59]]]

Note that when, for example, (x0, x1, x2) = (4,1,2) we have B[x0,x1,x2] = A[x2, x1, x0] = A[2,1,4] = 49.

If instead order = [0, 0, 0] and A as above, then we would have the output B be the 1-dimensional size-3 array B = [0, 26, 52] so that B[1] = B[x0] = A[x0,x0,x0] = A[1,1,1] = 26

Input

Here we use 0-indexing, but you may also use 1-indexing as is the APL default.

• A multidimensional or nested array A, of dimension n ≥ 1.

• A list order of n positive integers consisting of the integers {0,1,...,k} (or {1,...,k+1} for 1-index) for some k < n, in any order, possibly with repetitions.

Output

• A multidimensional or nested array representing the result of applying the dyadic transpose with those arguments. (The output will have dimension k+1.)

You may write a full program, function, etc. as allowed by the current standard on meta.

If your language has a builtin, it is encouraged to also write a solution without the builtin for the sake of an interesting answer.

Test cases

Test cases on TIO

Reference Python implementation coming soon.

Note for reading test cases: in APL, the penultimate and ultimate axes of an array are along columns and rows in that order.

Answer 1

APL (Dyalog Unicode), 34 bytes^SBCS

This is my colleague's code (slightly modified from Roger Hui: A History of APL in 50 Functions, chapter 30), posted with explicit permission.

Anonymous tacit infix lambda (can be used as drop-in for ⍉).

{⍵[(⊂⊂⍺)⌷¨,¨⍳⍺[⍋⍺]{⌊/⍵}⌸(⍴⍵)[⍋⍺]]}

Try it online!

{…} dfn; ⍺ is left argument (axes), ⍵ is right argument (array)
E.g. [2,2,1] and [[[1,2],[3,4]]]

 ⍵[…] index the array with:

  (⍴⍵)[…] the shape (lengths of axes) of the array, indexed by:
  [1,2,2]

   ⍋⍺ the grading vector (the indices that would sort them) of the axes
   [3,1,2]
  [2,1,2]

  ⍺[⍋⍺]{…}⌸ use the sorted axes as keys to group that, and for each group:
  [1,2,2] → {"1":[2],"2":[1,2]}

   {⌊/⍵} find the lowest axis length
   {"1":2,"2":1} → [2,1]

  ⍳ generate the indices in a Cartesian system of those dimensions
  [[[1,1]],[[2,1]]]

  ,¨ ensure that the indices of each coordinate is a vector (would be scalar if ⍵ is a vector)
  [[[1,1]],[[2,1]]]

  (…)⌷¨ index into each of those with the following:

   ⊂⊂⍺ the axes (doubly enclosed; once for ⌷ to select those cells along the first and only axis, and once for ¨ to pair up each vector on the right with the entire set of axes on the left)
   2 1 2
  [[[1,1,1]],[[1,2,1]]]
 [[1],[3]]