Given two numpy arrays, how to split one into an array of lists based on the second
Question:
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?
Answers:
You can use the split method from numpy:
output = np.split(values, np.cumsum(lengths))[:-1]
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?
You can use the split method from numpy:
output = np.split(values, np.cumsum(lengths))[:-1]