Numpy function to help avoid the use of any loops or if-else statements

user_12256590 0
Viewed 74

Given two arrays:

import numpy as np
array1 = np.array([7, 2, 4, 1, 20], dtype = "int")
array2 = np.array([2, 4, 4, 3, 10], dtype = "int")

and WITHOUT the use of any loop or if-else statement; I am trying to create a third array that will take the value equal to the sum of the (corresponding) elements from array1 and array2 if the element from array1 is bigger than the element from array2. If they are equal, the new element should have a value equal to their product. If the element from array2 is bigger, then the new element should be the difference between the element from array2 and array1

I have tried to implement this using python list and if-else statement with loop, but would like to know how to implement with numpy methods.

My implementation:

array1 = [7, 2, 4, 1, 20]
array2 = [2, 4, 4, 3, 10]
array3 = []

for i, j in enumerate(array1):
    if j>array2[i]:
        sum = j + array2[i]
        array3.append(sum)

    elif j==array2[i]:
        product = j * array2[i]
        array3.append(product)

    else:
        sub = array2[i] - j
        array3.append(sub)
print("output: ",array3)

output: [9, 2, 16, 2, 30]

numpy arrays python
4 Answers

You can use np.select here:

>>> import numpy as np
>>> array1 = np.array([7, 2, 4, 1, 20], dtype = "int")
>>> array2 = np.array([2, 4, 4, 3, 10], dtype = "int")

Here is the help for np.select:

select(condlist, choicelist, default=0)
    Return an array drawn from elements in choicelist, depending on conditions.

    Parameters
    ----------
    condlist : list of bool ndarrays
        The list of conditions which determine from which array in `choicelist`
        the output elements are taken. When multiple conditions are satisfied,
        the first one encountered in `condlist` is used.
    choicelist : list of ndarrays
        The list of arrays from which the output elements are taken. It has
        to be of the same length as `condlist`.
    default : scalar, optional
        The element inserted in `output` when all conditions evaluate to False.

    Returns
    -------
    output : ndarray
        The output at position m is the m-th element of the array in
        `choicelist` where the m-th element of the corresponding array in
        `condlist` is True.

So, applied to your problem:

>>> np.select(
...     [array1 > array2, array1 == array2, array1 < array2],
...     [array1 + array2, array1*array2, array2 - array1]
... )
array([ 9,  2, 16,  2, 30])
>>>
user_50144550

You can use three mask arrays, like so:

>>> array3 = np.zeros(array1.shape, dtype=array1.dtype)
>>> a1_gt = array1 > array2   # for when element at array 1 is greater
>>> a2_gt = array1 < array2   # for when element at array 2 is greater
>>> a1_eq_a2 = array1 == array2   # for when elements at array 1 and array 2 are equal
>>> array3[a1_gt] = array1[a1_gt] + array2[a1_gt]
>>> array3[a2_gt] = array2[a2_gt] - array1[a2_gt]
>>> array3[a1_eq_a2] = array2[a1_eq_a2] * array1[a1_eq_a2]
>>> array3
array([ 9.,  2., 16.,  2., 30.])
user_62751030

Using numpy.select with a default value:

array1 = np.array([7, 2, 4, 1, 20], dtype = "int")
array2 = np.array([2, 4, 4, 3, 10], dtype = "int")

np.select([array1>array2, array1<array2],
          [array1+array2, array2-array1],
          default=array1*array2)

output: array([ 9, 2, 16, 2, 30])

user_163434640

I renamed your arrays to a and b

print((a>b)*(a+b)+(a==b)*(a*b)+(a<b)*(b-a))

direct comparison between arrays give you boolean reasults that you can interpret as 0 or 1. That means a simple multiplication can turn an element "on" or "off". So we can just piece everything together.

user_55246370
Related