Search code examples
pythonarraysnumpyflatten

Is it possible to flatten an array with two different major orders without using for loop?


In python, I would like to flatten an array that are 3 dimensional. Say shape (10, 15, 200). But the first two would need to be flattened with row-major order while the remaining flatten operation would be done in column-major order.

I can probably do this in a for loop by iterating over operations like slicing the array, flattening it individually, and storing it in the main array. For example, I would slice the main array to be (10, 15) do a row-major flatten --> store it to the main with (150, 200), and then do a column-major flatten operation.

I am not sure if this is the most efficient method of doing this. Is there a better way to do this using numpy calls?


Solution

  • First use .reshape with the default "C" ordering (row-major) on the first two dimensions and then use .flatten with the "F" ordering (column-major) for the final result.

    import numpy as np
    
    a = np.arange(3*4*5).reshape(3,4,5)
    b = a.reshape(-1, a.shape[-1]).flatten("F")
    

    Result:

    array([ 0,  5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55,  1,  6, 11, 16, 21,
           26, 31, 36, 41, 46, 51, 56,  2,  7, 12, 17, 22, 27, 32, 37, 42, 47,
           52, 57,  3,  8, 13, 18, 23, 28, 33, 38, 43, 48, 53, 58,  4,  9, 14,
           19, 24, 29, 34, 39, 44, 49, 54, 59])