python:在函数内扩展numpy数组

Question

我有一个具有以下签名的 python 函数:

def merge(segments, indexes):

其中，segments 是一个 n 维 numpy 数组，indexs 是一个一维 numpy 数组。现在，我想调用以下函数:

np.where((segments == indexes[0]) | (segments == indexes[1]) | 
          ... segments == indexes[n])

但是，我不确定如何在 python 中的 where() 函数调用中动态生成此条件。

Answer 1

由于您有许多或条件，因此您可以使用np.in1d()来检查segments的每个元素是否存在于索引:

np.where(np.in1d(segments, indexes).reshape(segments.shape))

请注意，in1d() 的输出是一个扁平数组，需要重新调整形状，以便 where() 将返回正确的索引。