如果满足特定条件，则在python函数中应用另一个公式

Question

我正在尝试在python中编写一个函数，该函数使用特定公式生成输出数组，但如果满足特定条件，则使用另一个公式。

我写的代码是

import numpy as np

def fun(a,b):
for i in range(len(b)):
    if a != b[i]:
        return a+b[i]
    elif a == b[i]:
        return a-b[i]

a=1
b=np.array([5,4,1,2,3])

fun(a,b)

结果我得到了

6

而不是，

array([6, 5, 0, 3, 4])

任何人都可以将我引向该任务的适当代码，我想将想法推广到更复杂的公式上。

Answer 1

您的方法中的问题在于，它return是在第一次迭代中进行的，而不是创建列表并追加到列表中，然后返回该列表。

使用列表理解（不幸的np.array是不适用于生成器表达式）：

def fun(a, b):
    return np.array([n - 1 if n == a else n + 1 for n in b])

但是，由于您使用的是numpy，因此最好使用它提供的工具。

在这种情况下np.where：

a = 1
b = np.array([5,4,1,2,3])
print(np.where(b == 1, b - 1, b + 1))

产出

[6 5 0 3 4]