how to split a list of numpy arrays based on two other lists

user_14125436 0
Viewed 104

I have a list of numpy arrays stored in a list and want to split each array into two ones. This is my input data:

all_data=[[np.array([[1., 2., 0],[4., 5., 0],[0., 0., 0]]), np.array([[8., 8., 0],[7., 7., 0],[3., 3., 0]]),\
           np.array([[0., 0., 0],[1., 1., 0],[-1., 1., 0],[-2., -2., 0]])],\
          [np.array([[3., 0., 0],[1., 2., 0],[1., 1., 0],[0., 0., 0]]), np.array([[8., 8., 0],[7., 7., 0]]),\
           np.array([[5., 10., 0],[18., 25., 0],[25., 0., 0],[-10., 9., 0]])]]

And these are the lists I want to to do the split based on them:

sep1=[1, 2, 2, 1, 3, 1]
sep2=[2, 2, 1, 1, 1, 3]

First list (all_data[0][0]) has two arrays. First array should be split to two ones. Length of first split equals the first value in first_sep (first_sep[0]) and length of second equals the first values in second_sep (second_sep[0]). For next array of this first list (all_data[0][1]) again I have two splits but the legnths are coming from first_sep[2] and second_sep[2]. Then, to split the first array of the second list (all_data[1][0]), I should use first_sep[1] and second_sep[1] as the lengths. And for the last array of the second list I split it based on first_sep[3] and second_sep[3]. In fact numbers in the first_sep and second_sep are ordered based on the arrays in the within the lists and not arrays in a list. First numbers of both belong to the first array of the first list. Second numbers belong to the first array of the second list. In case of having three lists, third numbers belong to it but here third ones belong to second array of the first list. I also may have more than two arrays in each list. finally I want the following list as my output:

split_data=[[np.array([[1., 2., 0.]]), np.array([[4., 5., 0.],[0., 0., 0.]]),\
            np.array([[8., 8., 0.], [7., 7., 0.]]), np.array([[3., 3., 0.]]),\
            np.array([[0.,0.,0],[1., 1., 0],[-1., 1., 0]]), np.array([[-2., -2., 0]])],\
           [np.array([[3., 0., 0.], [1., 2., 0.]]), np.array([[1., 1.],[0., 0., 0.]]),\
            np.array([[8., 8., 0.]]), np.array([[7., 7., 0.]]), np.array([[5., 10., 0]]),\
           np.array([[18., 25., 0],[25., 0., 0],[-10., 9., 0]])]]

In advance I do appreciate any help to solve this issue.

numpy arrays python split
1 Answers

You could use Numpy's split function to split the arrays as you want.

Although not very efficient/elegant (your arrays have different shapes so vectorizing this seems impossible), the following line should work; note you don't even need the sep2 list, as you are splitting the arrays without taking out any elements (at least in your example).

split_data = [np.split(all_data[i][j], [sep1[2*j+i]]) for i in range(2) for j in range(len(all_data[i]))]
user_83456280
Related