How to split the numpy array into separate arrays in python. The number is splits is given by the user and the splitting is based on index

Question

I want to split my numpy array into separate arrays. The separation must be based on the index. The split count is given by the user.

For example, The input array: my_array=[1,2,3,4,5,6,7,8,9,10]

If user gives split count =2, then, the split must be like

my_array1=[1,3,5,7,9]
my_array2=[2,4,6,8,10]

if user gives split count=3, then the output array must be

my_array1=[1,4,7,10]
my_array2=[2,5,8]
my_array3=[3,6,9]

could anyone please explain, I did for split count 2 using even odd concept

for i in range(len(outputarray)):
    if i%2==0:
        even_array.append(outputarray[i])  
    else:
        odd_array.append(outputarray[i])

I don't know how to do the split for variable counts like 3,4,5 based on the index.

Answer 1

You can use indexing by vector (aka fancy indexing) for it:

>>> a=np.array([1,2,3,4,5,6,7,8,9,10])
>>> n = 3
>>> [a[np.arange(i, len(a), n)] for i in range(n)]

[array([ 1,  4,  7, 10]), array([2, 5, 8]), array([3, 6, 9])]

Explanation

arange(i, len(a), n) generates an array of integers starting with i, spanning no longer than len(a) with step n. For example, for i = 0 it generates an array

 >>> np.arange(0, 10, 3)
 array([0, 3, 6, 9])

Now when you index an array with another array, you get the elements at the requested indices:

 >>> a[[0, 3, 6, 9]]
 array([ 1,  4,  7, 10])

These steps are repeated for i=1..2 resulting in the desired list of arrays.