I have two numpy arrays: one containing arbitrary values, and one containing integers larger than 1. The sum of the integers is equal to the length of the first array. Sample:

values = np.array(["a", "b", "c", "d", "e", "f", "g", "h"])
lengths = np.array([1, 3, 2, 2])
len(values) == sum(lengths) # True

I would like to split the first array according to the lengths of the second array, and end up with something like:

output = np.array([["a"], ["b", "c", "d"], ["e", "f"], ["g", "h"]], dtype=object)

It’s easy to iterate over the array with a Python loop, but it’s also slow when both lists are very large (hundreds of millions of elements). Is there a way to do this operation using native numpy operations, which presumably should be must faster?

>Solution :

You can use the split method from numpy:

output = np.split(values, np.cumsum(lengths))[:-1]