是否有NumPy函数返回数组中某物的第一个索引?

我知道Python列表有一个方法可以返回某个对象的第一个索引:

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

NumPy数组也有类似的东西吗?

949699 次浏览
小开

NumPy中有很多操作可以放在一起来完成这个任务。这将返回等于item的元素的下标:

numpy.nonzero(array - item)

然后你可以取列表的第一个元素来得到一个元素。

小开

要在任何标准上建立索引，你可以这样做:

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]
小开
最佳答案

是的，给定一个数组，array和一个值，item来搜索，你可以使用np.where作为:

itemindex = numpy.where(array == item)

结果是一个元组，首先是所有的行索引，然后是所有的列索引。

例如，如果一个数组是二维的，它包含你的项目在两个位置，那么

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

将等于你的项目，因此将是:

array[itemindex[0][1]][itemindex[1][1]]
小开

如果你想用它作为其他东西的索引，如果数组是可广播的，你可以使用布尔索引;不需要显式索引。要做到这一点，绝对最简单的方法是基于真值进行索引。

other_array[first_array == item]

任何布尔运算都可以:

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

非零方法也接受布尔值:

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

两个0分别表示索引元组(假设first_array是1D)和索引数组中的第一项。

小开

如果你需要第一次出现的只有一个值的索引，你可以使用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的子序列的开头:

[# EYZ0, 1, 1 # EYZ1, 2 # EYZ2, # EYZ3, # EYZ2, # EYZ3, 8]

因此，它与寻找每个值的第一个发生略有不同。在你的程序中，你可以使用t的排序版本来得到你想要的:

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

您还可以将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。

小开

从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检查所有元素是否匹配。如果在数组的前面有匹配，生成器表达式会更快。

小开

在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.nonzero的任何方法。

然而，np.argwhere也可以用多维数组处理优雅的(你需要手动将它转换为元组而且，它不是短路的)，但如果没有找到匹配，它会失败:

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

对于1D数组，我建议使用np.flatnonzero(array == value)[0]，它等价于np.nonzero(array == value)[0][0]和np.where(array == value)[0][0]，但避免了对单元素元组进行开箱的麻烦。

小开

l.index(x)返回最小的我，这样我是列表中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和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),)

时间比较

只是检查对于大型数组，使用迭代器的解决方案更快当搜索项位于数组的开头时(在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:快速找到第一个值索引

小开

numpy_indexed包(免责声明，我是它的作者)包含一个向量化的等效列表。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，并且有各种处理缺失值的方法。

小开

注意:这是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
小开

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

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")
小开

使用ndindex

样本数组

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


...[[1,4],
[2,3]]

创建一个空列表来存储索引和元素元组

 index_elements = []
for i in np.ndindex(arr.shape):
index_elements.append((arr[i],i))

将元组列表转换为字典

 index_elements = dict(index_elements)
键是元素，值是元素 索引-使用键访问索引

 index_elements[4]
输出 
  ... (0,1)
小开

找到了另一个循环解决方案:

new_array_of_indicies = []


for i in range(len(some_array)):
if some_array[i] == some_value:
new_array_of_indicies.append(i)
小开

numpy中内置了一种相当习惯的向量化方法。它使用np.argmax()函数的一个奇怪之处来完成这一点——如果有许多值匹配，它将返回第一个匹配的索引。诀窍在于，对于布尔值，将永远只有两个值:True(1)和False(0)。因此，返回的索引将是第一个True的索引。

对于所提供的简单示例，您可以看到它在以下情况下工作

>>> np.argmax(np.array([1,2,3]) == 2)
1

一个很好的例子是计算桶，例如用于分类。假设你有一个切点数组，你想要一个“桶”;它对应于数组中的每个元素。该算法是计算cuts的第一个索引，其中x < cuts(在用np.Infitnity填充cuts之后)。我可以使用广播广播比较，然后沿# eyz0 - broadcasts轴应用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中的每个值都属于一个连续的箱子，具有定义良好且易于指定的边缘情况行为。

小开

另一个之前没有提到的选项是bisect模块，它也适用于列表，但需要一个预先排序的列表/数组:

import bisect
import numpy as np
z = np.array([104,113,120,122,126,138])
bisect.bisect_left(z, 122)

收益率

3

Bisect还会在您要查找的数字在数组中不存在时返回一个结果，以便将该数字插入正确的位置。

小开

8种方法的比较

TL;博士:

(注:适用于100M元素以下的1d数组)

  1. 为了获得最佳性能，使用index_of__v5 (numba + numpy.enumerate + for循环;参见下面的代码)。
  2. 如果numba不可用:
    1. 如果期望在前100k元素中找到目标值，则使用index_of__v7 (for循环+ enumerate)。
    2. 否则使用index_of__v2/v3/v4 (numpy.argmax或numpy.flatnonzero为基础)。

perf-plot

由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

小开

对于我的用例，我不能提前对数组排序，因为元素的顺序很重要。这是我的全部numpy实现:

import numpy as np


# The array in question
arr = np.array([1,2,1,2,1,5,5,3,5,9])


# Find all of the present values
vals=np.unique(arr)
# Make all indices up-to and including the desired index positive
cum_sum=np.cumsum(arr==vals.reshape(-1,1),axis=1)
# Add zeros to account for the n-1 shape of diff and the all-positive array of the first index
bl_mask=np.concatenate([np.zeros((cum_sum.shape[0],1)),cum_sum],axis=1)>=1
# The desired indices
idx=np.where(np.diff(bl_mask))[1]


# Show results
print(list(zip(vals,idx)))


>>> [(1, 0), (2, 1), (3, 7), (5, 5), (9, 9)]

我认为它解释了重复值的无序数组。

小开 
index_lst_form_numpy = pd.DataFrame(df).reset_index()["index"].tolist()

提交答案