python - numpy.where 与二维数组

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