NumPy Argument Array Split

Question

The function np.array_split(x, n) splits the array x into n roughly equally sized chunks. I am wondering what the most convenient form of this if one wants to obtain the indices of where the array is to be split. So an array of

ix = [(start1, end1), (start2, end2), ... (startn, endn)]

such that

np.array_split(x, n)[i] == x[ix[i][0]:ix[i][1]]

I can think of a few awkward ways of obtaining this but nothing simple.

Answer 1

You know the lengths of the sub arrays. Just use them to find start and end indices:

a = np.arange(10)
res = np.array_split(a, 3)
end = list(np.add.accumulate([len(x) for x in res]))
start = [0] + end[:-1]
ix = list(zip(start, end))

Now, the indices are:

>>> ix
[(0, 4), (4, 7), (7, 10)]

for this result:

>>> res
[array([0, 1, 2, 3]), array([4, 5, 6]), array([7, 8, 9])]

or:

for i in range(3):
    assert np.all(np.array_split(a, 3)[i] == a[ix[i][0]:ix[i][1]])