How to split a numpy array based on a tuple content?

Question

Let's say I've got an array [0, 1, 2, 3, 4, 5, 6, 7] and a tuple: (3, 3, 2).

I'm looking for a way to split my array to 3 array based on my tuple data:

[0, 1, 2]
[3, 4, 5]
[6, 7]

I can write a simple code like this to get what I want, however I'm looking for a correct and pythonic way to do this:

I used lists for simplicity.

a = [0, 1, 2, 3, 4, 5, 6, 7]
b = (3, 3, 2)

pointer = 0
for i in b:
        lst = []
        for j in range(i):
                lst.append(a[pointer])
                pointer += 1
        print(lst)

Or this one:

a = [0, 1, 2, 3, 4, 5, 6, 7]
b = (3, 3, 2)
pointer = 0
for i in b:
        lst = a[pointer:pointer+i]
        pointer += i
        print(lst)

Results:

[0, 1, 2]
[3, 4, 5]
[6, 7]

Usually it's a nuisance when posters ask about an array, but provide a list example. Here it's merely an annoyance. That said, `np.array_split` operates much like your 2nd solution, useing slices. But it expects a cumulative tuple. — hpaulj, Mar 24 '19 at 15:32

Querenker · Answer 1 · 2019-03-24T15:30:08.413

If you don't want to modify your input list, you can use an iterator and the itertools module.

>>> from itertools import islice
>>> a = [0, 1, 2, 3, 4, 5, 6, 7]
>>> b = (3, 3, 2)
>>> i = iter(a)
>>> [list(islice(i, x)) for x in b]
[[0, 1, 2], [3, 4, 5], [6, 7]]

In the first step you create an iterator, which starts at the first element of a. Then you iterate in a list comprehension over your numbers in b and in each step you pull accordingly many elements from the iterator and store them in your result list.

score 2 · Accepted Answer · answered Mar 24 '19 at 15:30

2

you can use the split method of numpy

import numpy as np

a = [0, 1, 2, 3, 4, 5, 6, 7]
b = (3, 3, 2)

c = np.split(a, np.cumsum(b)[:-1])

for r in c:
    print(r)

np.split(a, b) splits a by the indices in b along a given axis(0 by default).

answered Mar 24 '19 at 15:30

Yohai Magan

279
1
12

atru · Answer 3 · 2019-03-24T15:23:45.023

One simpler way is this:

a = [0, 1, 2, 3, 4, 5, 6, 7]
b = (3, 3, 2)

for ind in b:
    print(a[:ind])
    a = a[ind:]

It loops through slice sizes in b while shortening the original array every time. You can easily append the resulting slices as sublists if you need them for something else. It's almost like one of your solutions except it doesn't use any extra variables and iterates directly through elements of b.

Also, I wouldn't call variables a and b - surely not in this case where variables have clear meanings that you can express through their names. More meaningful names lessen bugs number and make code more clear, becomes a real difference with larger/more complex code. I'd call a at least in_list and b slices, but with more context this could be better.

score 0 · Answer 4 · answered Mar 24 '19 at 15:24

0

The most "concise" syntax would be :

ex_array = [0, 1, 2, 3, 4, 5, 6, 7] extuple = (3, 3, 2)

result = [ex_array[sum(extuple[:iii]):sum(extuple[:iii])+extuple[iii]] for iii in range(len(extuple))]

result would be a list of the expected sub-lists

answered Mar 24 '19 at 15:24

MaximGi

543
3
12

score 0 · Answer 5 · answered Mar 24 '19 at 15:35

Re-using the pairwise function from Compare two adjacent elements in same list, you could also:

from itertools import accumulate
from more_itertools import pairwise

a = [0, 1, 2, 3, 4, 5, 6, 7]
b = (3, 3, 2)

[a[slice(*s)] for s in pairwise(accumulate((0,)+b))]

That begin said, the np.split answer is probably faster (and easier to read).

How to split a numpy array based on a tuple content?

5 Answers5