Is there a NumPy function to return the first index of something in an array?

前端 未结 13 1639
萌比男神i
萌比男神i 2020-11-22 05:55

I know there is a method for a Python list to return the first index of something:

>>> l = [1, 2, 3]
>>> l.index(2)
1

Is

13条回答
  •  花落未央
    2020-11-22 06:11

    An alternative to selecting the first element from np.where() is to use a generator expression together with enumerate, such as:

    >>> import numpy as np
    >>> x = np.arange(100)   # x = array([0, 1, 2, 3, ... 99])
    >>> next(i for i, x_i in enumerate(x) if x_i == 2)
    2
    

    For a two dimensional array one would do:

    >>> x = np.arange(100).reshape(10,10)   # x = array([[0, 1, 2,... 9], [10,..19],])
    >>> next((i,j) for i, x_i in enumerate(x) 
    ...            for j, x_ij in enumerate(x_i) if x_ij == 2)
    (0, 2)
    

    The advantage of this approach is that it stops checking the elements of the array after the first match is found, whereas np.where checks all elements for a match. A generator expression would be faster if there's match early in the array.

提交回复
热议问题