Numpy Array non-sequentially divide the columns of the main array into n sub-arrays

I’ve been trying to do something like a numpy.array_split(), but to split it like this instead:

So It would return an array (for example let’s call it output[] ) with n 2D subarrays inside of it.

For example (for n = 3):

• output[0] would return the (yellow) subarray with columns a1, a4, a7, a10,
• output[1] would return the (red) subarray with columns a2, a5, a8,
• output[2] would return the (blue) subarray with columns a3, a6, a9.

def split(arr, n):
    output= [[] for _ in range(n)]
    for v, help in zip(arr, cycle(out)):
        help.append(v)
    return output

I don’t know how to combine rows into one 2D array, so I have many 1D arrays instead of one 2D.

>Solution :

Not sure if a native solution exists but you can use:

# get groups
group = np.arange(a.shape[1])%n
# groups sorting order
order = np.argsort(group)
# get counts of each group (in order as the output is sorted)
_, idx = np.unique(group, return_counts=True)
# split the reindexed array
out = np.split(a[:, order], np.cumsum(idx[:-1]), axis=1)

Output:

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

Used input:

array([[ 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]])