numpy - 如何按列比较二维数组和一维数组？

Question

我有两个 numpy 数组。其中一个是 2D，而另一个是 1D。

>>> a = np.arange(0,20).reshape(2,10)
>>> a
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]])
>>> b = np.full( a.shape[1], 10 )
>>> b
array([10, 10, 10, 10, 10, 10, 10, 10, 10, 10])

我想按列比较它们：

1. 如果其中的列元素a与 b 的列元素相同，则存储 a 的行号。
2. 否则，找到最接近的匹配项a并b存储行号。

在我的示例中，比较的输出应该是：

[ 1, [0,1], [0,1], [0,1], [0,1], [0,1], [0,1], [0,1], [0,1], [0,1] ]

我如何在 NumPy 中做到这一点？

我在想用np.where( a==b, run a function to get row(s) if same, run another function to get row(s) of diff )? 是这样吗？

标签： numpy