How can I reorder levels of a list in a performant way?

Question

I have the following problem and I am not sure how to speed it up:

I construct a list element by element like so:

a[[i]]=Table[some_function[j,k],{j,1,3},{k,1,6}]

This results in a list a with three indices and I would like to reorder its indices as follows:

Table[a[[i,j,k]],{j,1,3},{k,1,6},{i,1,4}]

Notice that initially the first index in a was i and now it is moved to the last position.

What I do at the moment is traverse the list with Table and reorder it elements as above, but I am wondering if there is a faster way to do this.

If this was a simple matrix, the Transpose operation would do the trick but I have no idea what the Transpose equivalent can apply here!

P.S.: the values of the list do not matter, just the order of its elements I need changing.

Answer 1

ClearAll[a]
a = Array[x, {3, 2, 4}];

MatrixForm[a, TableDirections -> {Column, Row, Row}]

flattened = Flatten[a, {{2}, {3}, {1}}]

MatrixForm[flattened, TableDirections -> {Column, Row, Row}]

Answer 2

Transpose generalizes nicely to arbitrary levels! Try

Transpose[a, {3, 1, 2}]

The specification of {3, 1, 2} causes the first level in a (corresponding to the first position in the list {3, 2, 1}) to become the third level in the result, the second level to become the first, etc.; see the docs for more examples.

Answer 3

Since the question is about performance, I think some timings should be included. There are two types of arrays and performance differs. Packed arrays are arrays whose elements are machine numbers of the same type. Functions are optimized to leverage strengths of the CPU on such arrays. Unpacked arrays are essentially, I think, lists of pointers (to lists of pointers...) to expressions.

Transpose cannot be beat on packed arrays, I think:

aa = RandomReal[1, {300, 200, 400}];
flattened = Flatten[aa, {{2}, {3}, {1}}]; // RepeatedTiming
transposed = Transpose[aa, {3, 1, 2}]; // RepeatedTiming
(*
{0.431257, Null}
{0.0949386, Null}
*)

Flatten seems faster on unpacked arrays, but Transpose plus repacking the array is faster:

bb = Developer`FromPackedArray@aa;
flattened = Flatten[bb, {{2}, {3}, {1}}]; // RepeatedTiming
transposed = Transpose[bb, {3, 1, 2}]; // RepeatedTiming
repacked = Transpose[Developer`ToPackedArray@bb,
            {3, 1, 2}]; // RepeatedTiming
(*
{0.792751, Null}
{1.13911, Null}
{0.313727, Null}
*)

But Transpose barely (but consistently) beats Flatten on this unpacked array:

cc = Array[x, {50, 40, 100}];
flattened = Flatten[cc, {{2}, {3}, {1}}]; // RepeatedTiming
transposed = Transpose[cc, {3, 1, 2}]; // RepeatedTiming
(*
{0.0179923, Null}
{0.0153114, Null}
*)