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

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

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


当前回答

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,并且有各种处理缺失值的方法。

其他回答

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

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

收益率

3

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

只是添加一个非常高性能和方便的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数组转换为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。

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

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,并且有各种处理缺失值的方法。