numpy.where with two-dimensional array

Question

可以使用numpy.where根据条件从两个数组中选择值：

import numpy
a = numpy.random.rand(5)
b = numpy.random.rand(5)
c = numpy.where(a > 0.5, a, b)  # okay

但是，如果数组具有更多维度，则不再有效：

import numpy
a = numpy.random.rand(5, 2)
b = numpy.random.rand(5, 2)
c = numpy.where(a[:, 0] > 0.5, a, b)  # !

Traceback (most recent call last):
File "p.py", line 10, in <module>
c = numpy.where(a[:, 0] > 0.5, a, b)  # okay
File "<__array_function__ internals>", line 6, in where
ValueError: operands could not be broadcast together with shapes (5,) (5,2) (5,2) 

我本来期望一个形状(5,2)的 numpy 数组.

这里有什么问题？如何解决它？

Answer 1

请记住，numpy广播只能从右侧进行，因此虽然(5,)形状的阵列可以使用(2,5)形状的阵列进行广播，但它们不能使用(5,2)形状的阵列进行广播。 要使用(5,2)形状的数组进行广播，您需要保持第二维，以便形状(5,1)(任何东西都可以用1广播(

因此，在为第二个维度编制索引时，您需要维护它(否则，当只有一个值存在时，它会删除索引维度(。 您可以通过将索引放在单元素列表中来执行此操作：

a = numpy.random.rand(5, 2)
b = numpy.random.rand(5, 2)
c = numpy.where(a[:, [0]] > 0.5, a, b) # works

Answer 2

您可以使用c = numpy.where(a > 0.5, a, b)

但是，如果您只想使用a的第一列，则需要考虑输出的形状。

我们先来看看这个操作的形状是什么

(a[:, 0] > 0.5).shape# 输出 (5，(

它是一维的

而 A 和 B 的形状为 (5， 2(

它是二维的，因此你不能广播这个

解决方案是将掩模操作重塑为形状 (5， 1(

您的代码应如下所示

a = numpy.random.rand(5, 2)
b = numpy.random.rand(5, 2)
c = numpy.where((a[:, 0] > 0.5).reshape(-1, 1), a, b)  # !

Answer 3

你可以试试：

import numpy
a = numpy.random.rand(5, 2)
b = numpy.random.rand(5, 2)
c = numpy.where(a > 0.5, a, b)

Answer 4

而不是：c = np.where(a>0.5,a,b)

您可以使用：c = np.array([a,b])[a>0.5]

如果 a 和 b 具有相同的形状，则适用于多维数组。