CODEWITHSUNDEEP
numpy
Borys
Borys

Reputation: 1423

Numpy get elements by incrementing the index of element

import numpy as np

a = np.array([1,2,13,7,3,5])
index = a == 13
print a[index+1]

The expected result is 7. How to get?

Upvotes: 0

Views: 94

Answers (4)

cel
cel

Reputation: 31349

I would do the following:

import numpy as np
a = np.array([1,2,13,7,3,5])
indices, = np.where(a[:-1] == 13)
print a[indices[0] + 1]

In python3 you can also do:

import numpy as np
a = np.array([1,2,13,7,3,5])
indices, = np.where(a[:-1] == 13)
first_index, *_ = indices
print(a[first_index + 1])

Which in my opinion is more readable. Unfortunately, extended tuple unpacking is not supported in python2 versions.

Both versions require that the problem is well defined and there is indeed a 13 in the list and that it has a predecessor. If this does not hold for your problem, the length of indices can be checked and in case of length 0 you may want to terminate with an appropriate error.

Upvotes: 1

chthonicdaemon
chthonicdaemon

Reputation: 19800

I think this is the most direct way:

import numpy as np

a = np.array([1,2,13,7,3,5])
print a[np.flatnonzero(a == 13) + 1]

Upvotes: 0

gboffi
gboffi

Reputation: 25053

Foreword

My proposed solutions are a generalization of the OP question, because the OP array had just an occurrence of the number 13 and here I treat the possibility of having more than a single occurrence.

All my solutions except one are based on the idea of the circularity of a list (as is made more clear by the name rotate that I've used for my helper function).

If there are use cases (I think there are such use cases)! where the assumption of circularity is not welcome, then use either my third piece of code or, better, the answer from @cel provided that you test the index before accessing the array or include the access in a try except clause.

Code Chunks

>>> import numpy as np
>>> a = np.array([1,2,13,7,3,5])
>>> i = a == 13
>>> print a[np.concatenate(([False],i[:-1]))]
[7]
>>>

another possibility

>>> print a[np.concatenate((i[-1:],i[:-1]))]
[7]
>>>

another one (fragile, it can index beyond the length of a, don't use this last one in this bare format, see foreword)

>>> print a[np.arange(len(a))[i]+1]
[7]
>>>

You have now an array of possibilities...

Elegance?

>>> def rotate(a,n):
...     n = n%len(a) 
...     return np.concatenate((a[-n:],a[:-n]))
...
>>> print a[rotate(i,1)]
[7]
>>>

Upvotes: 1

immerrr
immerrr

Reputation: 1273

The easiest way is to use nonzero. It returns a tuple of arrays, so you have to do some extra indexing:

In [1]: a = np.array([1,2,13,7,3,5])

In [2]: i = a == 13

In [3]: i
Out[3]: array([False, False,  True, False, False, False], dtype=bool)

In [4]: i.nonzero()
Out[4]: (array([2]),)

In [5]: i.nonzero()[0]
Out[5]: array([2])

In [6]: i.nonzero()[0][0]
Out[6]: 2

In [7]: a[i.nonzero()[0][0] + 1]
Out[7]: 7

Upvotes: 1

Related Questions