Elementwise, если функция elif в python с использованием массивов

У меня есть определение

def myfunc(a, b):
    if a < (b*10):
        result = a*2
    else:
        result = a*(-1)
    return result

Теперь это работает отлично, когда я кормлю свои значения a и b один за другим, используя для циклов, однако это берет навсегда (я упростил определение немного), и я знаю по опыту, что передача в значениях, поскольку массив ускорит его.

Итак, как мне изменить этот код, чтобы принимать массивы. Я использовал команды any() и all(), но я должен использовать их неправильно, поскольку моя функция только выплескивает одно значение, а не массив значений.

Примером моего желаемого результата будет:

>>>a = np.array([1,5,50,500])
>>>b = 1
>>>print myfunc(a, b)
array([-1, -5, 100, 1000])

Ответ 1

Вы можете использовать np.where:

def myfunc(a, b):
    return np.where(a < b*10, a*2, -a)

Например,

In [48]: a = np.array([1, 5, 50, 500])

In [49]: b = 1

In [50]: myfunc(a, b)
Out[50]: array([   2,   10,  -50, -500])

Обратите внимание, что вывод не совпадает с вашим желаемым выводом, но согласуется с кодом, который вы опубликовали. Разумеется, вы можете получить желаемый результат, изменив неравенство:

def myfunc(a, b):
    return np.where(a > b*10, a*2, -a)

затем

In [52]: myfunc(a, b)
Out[52]: array([  -1,   -5,  100, 1000])

Ответ 2

Используйте понимание списка:

myarray = [1, 5, 50, 500]
b = 1
[myfunc(a, b) for a in myarray]

Ответ 3

Ваша функция достаточно проста, чтобы ее можно было полностью отбросить:

arr = [1, 5, 50, 500]
arr = [a * 2 if a < b * 10 else -a for a in arr]