CODEWITHSUNDEEP
pythonnumpy
Zanam
Zanam

Reputation: 4807

Python: Creating list of subarrays

I have a massive array but for illustration I am using an array of size 14. I have another list which contains 2, 3, 3, 6. How do I efficiently without for look create a list of new arrays such that:

import numpy as np
A = np.array([1,2,4,5,7,1,2,4,5,7,2,8,12,3]) # array with 1 axis
subArraysizes = np.array( 2, 3, 3, 6 ) #sums to number of elements in A
B = list()
B[0] = [1,2]
B[1] = [4,5,7]
B[2] = [1,2,4]
B[3] = [5,7,2,8,12,3]

i.e. select first 2 elements from A store it in B, select next 3 elements of A store it in B and so on in the order it appears in A.

Upvotes: 2

Views: 1099

Answers (1)

Divakar
Divakar

Reputation: 221554

You can use np.split -

B = np.split(A,subArraysizes.cumsum())[:-1]

Sample run -

In [75]: A
Out[75]: array([ 1,  2,  4,  5,  7,  1,  2,  4,  5,  7,  2,  8, 12,  3])

In [76]: subArraysizes
Out[76]: array([2, 3, 3, 6])

In [77]: np.split(A,subArraysizes.cumsum())[:-1]
Out[77]: 
[array([1, 2]),
 array([4, 5, 7]),
 array([1, 2, 4]),
 array([ 5,  7,  2,  8, 12,  3])]

Upvotes: 3

Related Questions