python - Python:在函数内部扩展一个numpy数组

我有一个带有以下签名的python函数：

def merge(segments, indexes):

其中segments是n-d numpy数组，而indexs是一维numpy数组。现在，我要调用以下函数：

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

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

最佳答案

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

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

请注意，in1d()的输出是一个展平的数组，需要调整其形状，以使where()返回正确的索引。