numpy 广播和条件

Question

我确信这个问题已经在某个地方提出过，但我没有正确的关键字来找到解决方案...

我的问题是改进以下代码：

I = np.array([True,False])
x = np.array([1,2])
result = f(x) * (1 - I) + g(x) * I

实际上，I 被用作条件：如果 I 为 True，则执行第一件事，否则执行第二件事。 f 和 g 假设是 (2,) -> (2,) 函数。

有没有一种方法可以让我编写这样的代码，使 f 和 g 两个函数仅在需要时才计算，而不是在任何地方计算？这里我有 2 个值，但在应用程序中我有数百个并且这些函数很昂贵。

Answer 1

np.piecewise 可能是你想要的。例如：

I = np.array([True,False])
x = np.array([1,2])
f,g = np.square,np.negative
np.piecewise(x, I, [g,f])
# array([-1,  4])

一个潜在的陷阱：np.piecewise 的输出与x 的类型相同；所以如果x 是int 但产生float 值，这些值将被截断。为了避免这种将X 转换为float 或任何适当的方式。