我知道Python列表有一个方法可以返回某个对象的第一个索引:
>>> xs = [1, 2, 3]
>>> xs.index(2)
1
NumPy数组也有类似的东西吗?
我知道Python列表有一个方法可以返回某个对象的第一个索引:
>>> xs = [1, 2, 3]
>>> xs.index(2)
1
NumPy数组也有类似的东西吗?
当前回答
numpy中内置了一种相当习惯的向量化方法。它使用np.argmax()函数的一个奇怪之处来完成这一点——如果有许多值匹配,它将返回第一个匹配的索引。诀窍在于,对于布尔值,将永远只有两个值:True(1)和False(0)。因此,返回的索引将是第一个True的索引。
对于所提供的简单示例,您可以看到它在以下情况下工作
>>> np.argmax(np.array([1,2,3]) == 2)
1
一个很好的例子是计算桶,例如用于分类。假设你有一个切割点数组,你想要对应数组中每个元素的“桶”。该算法是计算x < cuts处的第一个切割索引(在使用np. infinity填充切割之后)。我可以使用broadcast来广播比较,然后沿着cuts-broadcast轴应用argmax。
>>> cuts = np.array([10, 50, 100])
>>> cuts_pad = np.array([*cuts, np.Infinity])
>>> x = np.array([7, 11, 80, 443])
>>> bins = np.argmax( x[:, np.newaxis] < cuts_pad[np.newaxis, :], axis = 1)
>>> print(bins)
[0, 1, 2, 3]
正如预期的那样,x中的每个值都属于一个连续的箱子,具有定义良好且易于指定的边界情况行为。
其他回答
您还可以将NumPy数组转换为list in - air并获取其索引。例如,
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。
只是添加一个非常高性能和方便的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,)
numpy_indexed包(免责声明,我是它的作者)包含一个向量化的等效list。ndarray的索引;那就是:
sequence_of_arrays = [[0, 1], [1, 2], [-5, 0]]
arrays_to_query = [[-5, 0], [1, 0]]
import numpy_indexed as npi
idx = npi.indices(sequence_of_arrays, arrays_to_query, missing=-1)
print(idx) # [2, -1]
这个解决方案具有向量化的性能,可以推广到ndarray,并且有各种处理缺失值的方法。
如果你只需要第一次出现一个值的索引,你可以使用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]),)
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