CODEWITHSUNDEEP
pythonarraysnumpy
SLhark
SLhark

Reputation: 177

Indexing in NumPy: Access every other group of values

The [::n] indexing option in numpy provides a very useful way to index every nth item in a list. However, is it possible to use this feature to extract multiple values, e.g. every other pair of values?

For example:

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

And I want to extract every other pair of values i.e. I want to return

a[0, 1, 4, 5, 8, 9,]

Of course the index could be built using loops or something, but I wonder if there's a faster way to use ::-style indexing in numpy but also specifying the width of the pattern to take every nth iteration of.

Thanks

Upvotes: 3

Views: 1338

Answers (2)

Dave Canton
Dave Canton

Reputation: 178

You can do that reshaping the array into a nx3 matrix, then slice up the first two elements for each row and finally flatten up the reshaped array:

a.reshape((-1,3))[:,:2].flatten()

resulting in:

array([ 0,  1,  3,  4,  6,  7,  9, 10])

Upvotes: 0

Divakar
Divakar

Reputation: 221704

With length of array being a multiple of the window size -

In [29]: W = 2 # window-size

In [30]: a.reshape(-1,W)[::2].ravel()
Out[30]: array([0, 1, 4, 5, 8, 9])

Explanation with breaking-down-the-steps -

# Reshape to split into W-sized groups
In [43]: a.reshape(-1,W)
Out[43]: 
array([[ 0,  1],
       [ 2,  3],
       [ 4,  5],
       [ 6,  7],
       [ 8,  9],
       [10, 11]])

# Use stepsize to select every other pair starting from the first one
In [44]: a.reshape(-1,W)[::2]
Out[44]: 
array([[0, 1],
       [4, 5],
       [8, 9]])

# Flatten for desired output
In [45]: a.reshape(-1,W)[::2].ravel()
Out[45]: array([0, 1, 4, 5, 8, 9])

If you are okay with 2D output, skip the last step as that still be a view into the input and virtually free on runtime. Let's verify the view-part -

In [47]: np.shares_memory(a,a.reshape(-1,W)[::2])
Out[47]: True

For generic case of not necessarily a multiple, we can use a masking based one -

In [64]: a[(np.arange(len(a))%(2*W))<W]
Out[64]: array([0, 1, 4, 5, 8, 9])

Upvotes: 7

Related Questions