'wrap around' when slicing in python/numpy

Viewed 190

I have a NumPy array

import numpy as np
a = np.array([10, 20, 30, 40])

and I'd like to slice it in such a way, that the entire array is returned, with the first element appended to the end, so array([10, 20, 30, 40, 10]), or (as a normal list) [10, 20, 30, 40, 10].

4 Answers
>>> a = np.array([10, 20, 30, 40])
>>> np.take(a, np.arange(5), mode='wrap')
array([10, 20, 30, 40, 10])

Consider cycle and islice from the itertools library to achieve this

>>> import numpy as np
>>> import itertools
>>> a = np.array([10, 20, 30, 40])
>>> list(itertools.islice(itertools.cycle(a), len(a) + 1))
[10, 20, 30, 40, 10]

Note that this will produce an iterable, which you do not need to consume at once.

There are several ways to do this, as collected from my own findings, from the comments below the original question, and from answers by others:

b1 = a.tolist() + [a[0]]
b2 = np.append(a, a[0])

c1 = [a[i%len(a)] for i in range(len(a)+1)]
c2 = a[np.arange(len(a)+1) % len(a)]

d1 = np.array((*a, a[0]))
d2 = [*a, a[0]]

e = np.resize(a, len(a)+1)

f = np.take(a, range(len(a)+1), mode='wrap')

from itertools import cycle, islice
g = list(islice(cycle(a), len(a)+1))

Of these, c1, c2, e, f and g can be most flexibly/easily adjusted in case more elements (than just the first) must be appended, by changing the len(a)+1 part of the assignment; e is my favourite for readability and brevity.

I too needed to 'wrap around' a numpy array. However I, like many who likely view this question, needed to wrap around a multidimensional array.

To do so, simply modify the currently accepted and favored answer.

>>> a = np.array([10, 20, 30, 40])
>>> wrap_by = 1
>>> np.resize(a, len(a)+wrap_by)
array([10, 20, 30, 40, 10])

>>> b = np.array([[10, 20], [30, 40], [50, 60]])
>>> wrap_by = 1
>>> assert len(b.shape)>1 # 'b' is indeed a multidimensional array
>>> np.resize(b, (len(b)+wrap_by, b.shape[1]))
array([[10, 20],
   [30, 40],
   [50, 60],
   [10, 20]])
Related