python - 根据多个条件设置numpy数组的值

Question

我试图将 numpy 数组中的值设置为零，如果它等于列表中的任何数字。

让我们考虑以下数组

  a = numpy.array([[1, 2, 3], [4, 8, 6], [7, 8, 9]])

我想将列表 [1, 2, 8] 中 a 的多个元素设置为 0。

结果应该是

   [[0, 0, 3],
    [4, 0, 6],
    [7, 0, 9]]

对于单个元素来说很简单

   a[a == 1] = 0

以上仅适用于单个整数。它如何适用于列表？

Answer 1

使用 np.in1d 您可以执行以下操作:

>>> a = np.array([[1, 2, 3], [4, 8, 6], [7, 8, 9]])
>>> np.in1d(a, [1, 2, 8])
array([ True,  True, False, False,  True, False, False,  True, False], dtype=bool)
>>> a[np.in1d(a, [1, 2, 8]).reshape(a.shape)] = 0
>>> a
array([[0, 0, 3],
       [4, 0, 6],
       [7, 0, 9]])