python - Numpy:检查多维数组中的元素是否在元组中

Question

看来我还在纠结the "in" operator in numpy .情况是这样的:

>>> a = np.random.randint(1, 10, (2, 2, 3))
>>> a
array([[[9, 8, 8],
        [4, 9, 1]],

       [[6, 6, 3],
        [9, 3, 5]]])

我想获取那些第二个元素在 (6, 8) 中的三元组的索引。我凭直觉尝试的方式是:

>>> a[:, :, 1] in (6, 8)
ValueError: The truth value of an array with more than one element...

我的最终目标是在那些位置插入乘以二的数字。使用上面的例子，a 应该变成:

array([[[9, 18, 8],   #8 @ pos #2 --> replaced by 9 @ pos #1 by 2
        [4, 9, 1]],

       [[6, 12, 3],   #6 @ pos #2 --> replaced by 6 @ pos #1 by 2
        [9, 3, 5]]])

提前感谢您的建议和时间!

Answer 1

这是一个适用于任意长度元组的方法。它使用 numpy.in1d 函数。

import numpy as np
np.random.seed(1)

a = np.random.randint(1, 10, (2, 2, 3))
print(a)

check_tuple = (6, 9, 1)

bool_array = np.in1d(a[:,:,1], check_tuple)
ind = np.where(bool_array)[0]
a0 = a[:,:,0].reshape((len(bool_array), ))
a1 = a[:,:,1].reshape((len(bool_array), ))
a1[ind] = a0[ind] * 2

print(a)

输出:

[[[6 9 6]
  [1 1 2]]

 [[8 7 3]
  [5 6 3]]]

[[[ 6 12  6]
  [ 1  2  2]]

 [[ 8  7  3]
  [ 5 10  3]]]