我知道Python列表有一个方法可以返回某个对象的第一个索引:
>>> xs = [1, 2, 3]
>>> xs.index(2)
1
NumPy数组也有类似的东西吗?
我知道Python列表有一个方法可以返回某个对象的第一个索引:
>>> xs = [1, 2, 3]
>>> xs.index(2)
1
NumPy数组也有类似的东西吗?
当前回答
8种方法的比较
TL; diana:
(注:适用于100M元素以下的1d数组)
为了获得最佳性能,请使用index_of__v5 (numba + numpy. 5)。枚举+ for循环;参见下面的代码)。 如果numba不可用: 如果期望在前100k个元素中找到目标值,请使用index_of__v7 (for循环+枚举)。 否则使用index_of__v2/v3/v4 (numpy. exe)。Argmax或numpy。基于flatnonzero)。
由perfplot提供
import numpy as np
from numba import njit
# Based on: numpy.argmax()
# Proposed by: John Haberstroh (https://stackoverflow.com/a/67497472/7204581)
def index_of__v1(arr: np.array, v):
is_v = (arr == v)
return is_v.argmax() if is_v.any() else -1
# Based on: numpy.argmax()
def index_of__v2(arr: np.array, v):
return (arr == v).argmax() if v in arr else -1
# Based on: numpy.flatnonzero()
# Proposed by: 1'' (https://stackoverflow.com/a/42049655/7204581)
def index_of__v3(arr: np.array, v):
idxs = np.flatnonzero(arr == v)
return idxs[0] if len(idxs) > 0 else -1
# Based on: numpy.argmax()
def index_of__v4(arr: np.array, v):
return np.r_[False, (arr == v)].argmax() - 1
# Based on: numba, for loop
# Proposed by: MSeifert (https://stackoverflow.com/a/41578614/7204581)
@njit
def index_of__v5(arr: np.array, v):
for idx, val in np.ndenumerate(arr):
if val == v:
return idx[0]
return -1
# Based on: numpy.ndenumerate(), for loop
def index_of__v6(arr: np.array, v):
return next((idx[0] for idx, val in np.ndenumerate(arr) if val == v), -1)
# Based on: enumerate(), for loop
# Proposed by: Noyer282 (https://stackoverflow.com/a/40426159/7204581)
def index_of__v7(arr: np.array, v):
return next((idx for idx, val in enumerate(arr) if val == v), -1)
# Based on: list.index()
# Proposed by: Hima (https://stackoverflow.com/a/23994923/7204581)
def index_of__v8(arr: np.array, v):
l = list(arr)
try:
return l.index(v)
except ValueError:
return -1
去Colab
其他回答
对于1D数组,我推荐np。平坦非零(array == value)[0],它等价于np。非零(array == value)[0][0]和np。其中(array == value)[0][0],但避免了对一个单元素元组开箱的丑陋。
只是添加一个非常高性能和方便的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的方法都要快得多(因为它使操作短路)。Where或np. non0。
然而np。Argwhere也可以优雅地处理多维数组(你需要手动将它转换为元组,而且不会短路),但如果没有找到匹配,它就会失败:
>>> tuple(np.argwhere(arr1 == 2)[0])
(2, 2, 2)
>>> tuple(np.argwhere(arr2 == 2)[0])
(5,)
从np.where()中选择第一个元素的替代方法是使用生成器表达式和enumerate,例如:
>>> 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
对于二维数组,可以这样做:
>>> 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)
这种方法的优点是,它在找到第一个匹配后停止检查数组的元素,而np。Where检查所有元素是否匹配。如果在数组的前面有匹配,生成器表达式会更快。
注意:这是python 2.7版本
您可以使用lambda函数来处理这个问题,它既适用于NumPy数组,也适用于列表。
your_list = [11, 22, 23, 44, 55]
result = filter(lambda x:your_list[x]>30, range(len(your_list)))
#result: [3, 4]
import numpy as np
your_numpy_array = np.array([11, 22, 23, 44, 55])
result = filter(lambda x:your_numpy_array [x]>30, range(len(your_list)))
#result: [3, 4]
你可以用
result[0]
获取筛选元素的第一个索引。
对于python 3.6,使用
list(result)
而不是
result
L.index (x)返回最小的I,使得I是x在列表中第一次出现的索引。
可以放心地假设,Python中的index()函数的实现使它在找到第一个匹配后停止,这将导致最佳的平均性能。
要在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和非零返回数组,因此需要选择第一个元素来获取索引。
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),)
时间比较
只是检查对于大型数组,当搜索项位于数组的开头时,使用迭代器的解决方案更快(在IPython shell中使用%timeit):
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:快速找到第一个值索引