python中帶有布爾值的向量方程

Question

我正在嘗試創建一個方程，您可以在其中輸入一個數組或向量，並且它包含一個布爾函數。 對於滿足此條件的數組中的項目 (bool=True)，然后以一種方式繼續求解方程以生成另一個數組。

我在這里附上了在 R 中工作的類似代碼，並希望在 python 中做類似的事情

a <- c(0,1,2,3,4,5)
b <- c(1,1,2,2,3,3)

a-b+5*(a==0|b==0)

其輸出是一個向量：

[1] 4 0 0 1 1 2

有誰知道如何在 python3 中做類似的事情，也許用 numpy？

Answer 1

與numpy ：

import numpy as np
a = np.array([0,1,2,3,4,5])
b = np.array([1,1,2,2,3,3])
a-b+5*((a==0)|(b==0))
#> array([4, 0, 0, 1, 1, 2])

Answer 2

方法很多很多。 這是一個：

a = range(6)
b = [1, 1, 2, 2, 3, 3]
r = [ai - bi + 5*(ai==0 or bi==0) for ai, bi in zip(a,b)]

Answer 3

一種使用numpy ：

import numpy as np
a = np.arange(6)
b = np.repeat(np.arange(1,4),2)
a-b+5*(np.equal(a,0) | np.equal(b,0))