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

前端 未结 13 1814
萌比男神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:16

    If you need the index of the first occurrence of only one value, you can use nonzero (or where, which amounts to the same thing in this case):

    >>> t = array([1, 1, 1, 2, 2, 3, 8, 3, 8, 8])
    >>> nonzero(t == 8)
    (array([6, 8, 9]),)
    >>> nonzero(t == 8)[0][0]
    6
    

    If you need the first index of each of many values, you could obviously do the same as above repeatedly, but there is a trick that may be faster. The following finds the indices of the first element of each subsequence:

    >>> nonzero(r_[1, diff(t)[:-1]])
    (array([0, 3, 5, 6, 7, 8]),)
    

    Notice that it finds the beginning of both subsequence of 3s and both subsequences of 8s:

    [1, 1, 1, 2, 2, 3, 8, 3, 8, 8]

    So it's slightly different than finding the first occurrence of each value. In your program, you may be able to work with a sorted version of t to get what you want:

    >>> st = sorted(t)
    >>> nonzero(r_[1, diff(st)[:-1]])
    (array([0, 3, 5, 7]),)
    

提交回复
热议问题