在numpy.select条件中使用

Question

我试图在numpy.select中使用“ in”。

x = np.arange(10)
condlist = [x in [2,3,4], x>5]
choicelist = [x, x**2]
np.select(condlist, choicelist)

有办法使它工作吗？

Answer 1

您应该改用isin：

x = np.arange(10)
condlist = [np.isin(x, [2,3,4]), x>5]
choicelist = [x, x**2]
np.select(condlist, choicelist)

输出：

array([ 0,  0,  2,  3,  4,  0, 36, 49, 64, 81])