Weena
Weena

Reputation: 55

How to calculate each single element of a numpy array based on conditions

I have a one column numpy array with multiple elements. I want to write a function that takes each element and evaluates if it meets a condition. Then depending on the result, it should use equations to calculate something. The calculation result should have the same size as the input array.

This is a simple example of what I want to do. (The actual code will be more complex). I know why it doesn't work but can't seem to find a solution.

import numpy as np
array1 = np.arange(1,11,1)

def test(array1):
    value1 = 20
    i = 0
    value3 = array1[i]
    while array1 > value3 and i < value1:
        i =+ 1
        value3 = array1[i]

test(array1)

I have tried to find a solution:

EDIT: For my complete solution both the for loop version by Zulfiqaar and the where function as shown by Siva Kowuru are possible. Since the where function seems faster and more convenient I marked it as the accepted answer.

Upvotes: 1

Views: 3572

Answers (3)

Siva Kovvuru
Siva Kovvuru

Reputation: 86

You could add the optional parameters on numpy.where(condition[, x, y]) to keep the same array size.

When True, yield x, otherwise yield y

https://docs.scipy.org/doc/numpy/reference/generated/numpy.where.html


To compare to the answer from @Zulfiqaar, the solution would be

np.where(array1 < 3, array1 ** 2, array1 / 3)

Note that the condition is just a bool array, so you can combine multiple conditions at once using binary operators. For instance comparing to another array:

array2 = np.random.randint(1, 11, 10)

>>> array2
Out[]: array([ 6,  1,  6,  2, 10,  4,  9, 10,  3,  3])

>>> np.where((array1 < 3) & (array1 > 5), array1 ** 2, array1 / 3)
Out[]:
array([ 0.33333333,  4.        ,  1.        ,  1.33333333,  1.66666667,
        2.        ,  2.33333333,  2.66666667,  3.        ,  3.33333333])

Upvotes: 3

Spas
Spas

Reputation: 860

You can use map with lambda function instead of looping through all the elements:

value3 = array1[0]
map_array = np.array(list(map(lambda x: x**2 if x < value3 else x/3, array1)))

If you aren't using python 3 you won't need to convert it to list from the map function.

Upvotes: 0

Zulfiqaar
Zulfiqaar

Reputation: 633

should use if for conditions:

import numpy as np

array1 = np.arange(1, 11, 1)
array2 = np.zeros(10)


for eachValue in range(len(array1)):                #iterate through array
    if (array1[eachValue] <3 ):                     #check for your condition
        array2[eachValue] = array1[eachValue] ** 2  #apply equation
    else:                                           #if condition not met
        array2[eachValue] = array1[eachValue] /  3  #do something else

Upvotes: 1

Related Questions