是否有NumPy函数返回数组中某些内容的第一个索引？

Question

我知道Python列表有一种方法可以返回第一个索引:

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

NumPy阵列有类似的东西吗？

Answer 1

是的,这里给出了一个NumPy数组的答案array,以及一个值item,用于搜索:

itemindex = numpy.where(array==item)

结果是一个元组,首先是所有行索引,然后是所有列索引.

例如,如果一个数组是两个维度,那么它在两个位置包含您的项目

array[itemindex[0][0]][itemindex[1][0]]

将等于你的项目,所以会

array[itemindex[0][1]][itemindex[1][1]]

numpy.where

Answer 2

如果你需要第一次出现的索引只有一个值,你可以使用nonzero(或者where,在这种情况下相同的东西):

>>> 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

如果您需要多个值中的每个值的第一个索引,您显然可以重复执行与上面相同的操作,但有一个技巧可能更快.以下查找每个子序列的第一个元素的索引:

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

请注意,它找到3s的子序列的开始和8s的两个子序列:

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

所以它与找到每个值的第一次出现略有不同.在您的程序中,您可以使用已排序的版本t来获得您想要的内容:

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

Answer 3

您还可以将NumPy数组转换为空中列表并获取其索引.例如,

l = [1,2,3,4,5] # Python list
a = numpy.array(l) # NumPy array
i = a.tolist().index(2) # i will return index of 2
print i

它将打印1.

Answer 4

只需添加一个非常高性能和方便的numba替代方案,np.ndenumerate以找到第一个索引:

from numba import njit
import numpy as np

@njit
def index(array, item):
    for idx, val in np.ndenumerate(array):
        if val == item:
            return idx
    # If no item was found return None, other return types might be a problem due to
    # numbas type inference.

这非常快,并且可以自然地处理多维数组:

>>> arr1 = np.ones((100, 100, 100))
>>> arr1[2, 2, 2] = 2

>>> index(arr1, 2)
(2, 2, 2)

>>> arr2 = np.ones(20)
>>> arr2[5] = 2

>>> index(arr2, 2)
(5,)

这可能比使用或使用任何方法快得多(因为它使操作短路).np.wherenp.nonzero

但是np.argwhere也可以优雅地处理多维数组(您需要手动将其转换为元组并且它不会被短路)但如果找不到匹配则会失败:

>>> tuple(np.argwhere(arr1 == 2)[0])
(2, 2, 2)
>>> tuple(np.argwhere(arr2 == 2)[0])
(5,)

Answer 5

如果您打算将其用作其他内容的索引,则可以使用布尔索引(如果数组是可广播的); 你不需要明确的指数.绝对最简单的方法是简单地根据真值进行索引.

other_array[first_array == item]

任何布尔操作都有效:

a = numpy.arange(100)
other_array[first_array > 50]

非零方法也采用布尔值:

index = numpy.nonzero(first_array == item)[0][0]

两个零用于索引的元组(假设first_array是1D),然后是索引数组中的第一个项.

Answer 6

对于一维排序数组，使用返回 NumPy 整数（位置）的numpy.searchsorted会更简单和有效 O(log(n) ) 。例如，

arr = np.array([1, 1, 1, 2, 3, 3, 4])
i = np.searchsorted(arr, 3)

只要确保数组已经排序

还要检查返回的索引 i 是否确实包含搜索的元素，因为 searchsorted 的主要目标是找到应该插入元素以保持顺序的索引。

if arr[i] == 3:
    print("present")
else:
    print("not present")

Answer 7

l.index(x)返回最小的i,使得i是列表中第一次出现x的索引.

可以安全地假设index()Python 中的函数已经实现,以便在找到第一个匹配项后停止,这样可以获得最佳的平均性能.

要查找在NumPy数组中第一次匹配后停止的元素,请使用迭代器(ndenumerate).

In [67]: l=range(100)

In [68]: l.index(2)
Out[68]: 2

NumPy数组:

In [69]: a = np.arange(100)

In [70]: next((idx for idx, val in np.ndenumerate(a) if val==2))
Out[70]: (2L,)

请注意,这两种方法index(),并next在未找到该元素,返回一个错误.使用时next,可以使用第二个参数在未找到元素的情况下返回特殊值,例如

In [77]: next((idx for idx, val in np.ndenumerate(a) if val==400),None)

NumPy(argmax,where和nonzero)中还有其他函数可用于查找数组中的元素,但它们都有缺点,即遍历整个数组查找所有出现的内容,因此未针对查找第一个元素进行优化.还要注意where并nonzero返回数组,因此需要选择第一个元素来获取索引.

In [71]: np.argmax(a==2)
Out[71]: 2

In [72]: np.where(a==2)
Out[72]: (array([2], dtype=int64),)

In [73]: np.nonzero(a==2)
Out[73]: (array([2], dtype=int64),)

时间比较

只检查大型数组,当搜索项位于数组的开头时(使用%timeitIPython shell),使用迭代器的解决方案会更快:

In [285]: a = np.arange(100000)

In [286]: %timeit next((idx for idx, val in np.ndenumerate(a) if val==0))
100000 loops, best of 3: 17.6 µs per loop

In [287]: %timeit np.argmax(a==0)
1000 loops, best of 3: 254 µs per loop

In [288]: %timeit np.where(a==0)[0][0]
1000 loops, best of 3: 314 µs per loop

这是一个开放的NumPy GitHub问题.

另见:Numpy:快速找到第一个价值指数

Answer 8

要对任何条件进行索引,您可以执行以下操作:

In [1]: from numpy import *
In [2]: x = arange(125).reshape((5,5,5))
In [3]: y = indices(x.shape)
In [4]: locs = y[:,x >= 120] # put whatever you want in place of x >= 120
In [5]: pts = hsplit(locs, len(locs[0]))
In [6]: for pt in pts:
   .....:         print(', '.join(str(p[0]) for p in pt))
4, 4, 0
4, 4, 1
4, 4, 2
4, 4, 3
4, 4, 4

这里有一个快速的函数来执行list.index()所做的事情,但如果没有找到则不会引发异常.注意 - 这在大型阵列上可能非常慢.如果你更愿意将它作为一种方法使用,你可以将其修补到数组上.

def ndindex(ndarray, item):
    if len(ndarray.shape) == 1:
        try:
            return [ndarray.tolist().index(item)]
        except:
            pass
    else:
        for i, subarray in enumerate(ndarray):
            try:
                return [i] + ndindex(subarray, item)
            except:
                pass

In [1]: ndindex(x, 103)
Out[1]: [4, 0, 3]

Answer 9

8种方法比较

\n

长话短说：

\n

（注：适用于100M元素以下的一维数组。）

\n

\n
1. 为了获得最大性能使用index_of__v5（numba\xc2\xa0+ numpy.enumerate+for循环；参见下面的代码）。
\n
2. 如果numba不可用：\n
   \n
   1. 如果希望在前 100k 个元素中找到目标值，请使用index_of__v7( for\xc2\xa0loop + )\xc2\xa0。enumerate
   \n
   2. 否则使用index_of__v2/v3/v4\xc2\xa0( numpy.argmax\xc2\xa0or numpy.flatnonzero\xc2\xa0based)。
   \n
   \n
\n

\n

\n

由perfplot提供支持

\n

import numpy as np\nfrom numba import njit\n\n# Based on: numpy.argmax()\n# Proposed by: John Haberstroh (https://stackoverflow.com/a/67497472/7204581)\ndef index_of__v1(arr: np.array, v):\n    is_v = (arr == v)\n    return is_v.argmax() if is_v.any() else -1\n\n\n# Based on: numpy.argmax()\ndef index_of__v2(arr: np.array, v):\n    return (arr == v).argmax() if v in arr else -1\n\n\n# Based on: numpy.flatnonzero()\n# Proposed by: 1\'\' (https://stackoverflow.com/a/42049655/7204581)\ndef index_of__v3(arr: np.array, v):\n    idxs = np.flatnonzero(arr == v)\n    return idxs[0] if len(idxs) > 0 else -1\n\n\n# Based on: numpy.argmax()\ndef index_of__v4(arr: np.array, v):\n    return np.r_[False, (arr == v)].argmax() - 1\n\n\n# Based on: numba, for loop\n# Proposed by: MSeifert (https://stackoverflow.com/a/41578614/7204581)\n@njit\ndef index_of__v5(arr: np.array, v):\n    for idx, val in np.ndenumerate(arr):\n        if val == v:\n            return idx[0]\n    return -1\n\n\n# Based on: numpy.ndenumerate(), for loop\ndef index_of__v6(arr: np.array, v):\n    return next((idx[0] for idx, val in np.ndenumerate(arr) if val == v), -1)\n\n\n# Based on: enumerate(), for loop\n# Proposed by: Noyer282 (https://stackoverflow.com/a/40426159/7204581)\ndef index_of__v7(arr: np.array, v):\n    return next((idx for idx, val in enumerate(arr) if val == v), -1)\n\n\n# Based on: list.index()\n# Proposed by: Hima (https://stackoverflow.com/a/23994923/7204581)\ndef index_of__v8(arr: np.array, v):\n    l = list(arr)\n    try:\n        return l.index(v)\n    except ValueError:\n        return -1\n

\n

前往科拉布

\n

Answer 10

对于1D阵列,我建议np.flatnonzero(array == value)[0],这相当于两者np.nonzero(array == value)[0][0],np.where(array == value)[0][0]但避免了取消装箱1元素元组的丑陋.