python get next x elements of numpy array remembering index of the last get

acmpo6ou 0
Viewed 27

I had a problem, and I solved it by creating a class:

class ArrayIter:
    def __init__(self, array):
        self.array = array
        self.index = 0

    def __getitem__(self, n):
        res = self.array[self.index: self.index + n]
        self.index += n
        return res

it can be used like this:

>>> a = np.arange(20)
>>> iter = ArrayIter(a)
>>> iter[3]
array([0, 1, 2])
>>> iter[3]
array([3, 4, 5])
>>> iter[6]
array([ 6,  7,  8,  9, 10, 11])

Now something tells me I have reinvented a wheel. Is there something in python's standard library that can do the same? Keep in mind that I'm using numpy arrays and not lists.

numpy arrays python
1 Answers

Try generators:

a = np.arange(20)
it = iter(a)

def get(n):
    for _ in range(n):
        yield next(it)

print(list(get(2)))
print(list(get(3)))
print(list(get(10)))
print(list(get(1)))

# [0, 1]
# [2, 3, 4]
# [5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
# [15]

Details: What does the "yield" keyword do?

zalevskiaa0
Related