python - 查找 numpy 数组中某个非 X 或 Y 的值第一次出现的索引

Question

我有一个以下形式的向量:

import numpy as np
vec = np.array([2, 2, 2, 51, 51, 52, 52, 14, 14, 14, 51, 51, 52, 52])

有没有一种 numpy-thonic 方法可以找到第一次出现的值不是(例如)51 或 52 的索引？换句话说，该函数将返回以下索引:[0, 7]，其中 0 是第一个出现 2 的索引，7 是第一个出现 14 的索引。

Answer 1

np.unique如果指定 return_index=True，则返回每个数字的第一个索引。您可以使用“np.isin”轻松过滤结果。 :

u, i =  np.unique(vec, return_index=True)
result = i[np.isin(u, [51, 52], invert=True)]

这样做的好处是，与原始数据相比，u 的搜索空间显着减少。与显式否定结果掩码相比，使用 invert=True 还可以加快速度。

np.isin 的一个版本依赖于数据已经排序这一事实，可以使用np.searchsorted 来制作。像这样:

def isin_sorted(a, i, invert=False):
    ind = np.searchsorted(a, i)
    ind = ind[a[ind.clip(max=a.size)] == i]
    if invert:
        mask = np.ones(a.size, dtype=bool)
        mask[ind] = False
    else:
        mask = np.zeros(a.size, dtype=bool)
        mask[ind] = True
    return mask

在调用 np.unique 之后，您可以使用此版本代替 np.isin，它始终返回一个排序数组。对于足够大的vec和排除列表，它会更有效:

result = i[isin_sorted(u, [51, 52], invert=True)]