CODEWITHSUNDEEP
pythonnumpy
okkhoy
okkhoy

Reputation: 1338

conditional product in numpy

I have a list which controls what terms in the data list have to be multiplied 

control_list = [1, 0, 1, 1, 0]
data_list = [5, 4, 5, 5, 4]

I need to find product of elements in the data_list for which the control_list has 1. My current attempt is naive and looks ugly!

product = 1
for i in range(len(control_list)):
    if control_list[i]: 
        product *= data_list[i]

I looked at numpy.where() to get the required elements in data_list but it looks like I did not get it right:

numpy.where(control_list, data_list)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-12-1534a6838544> in <module>()
----> 1 numpy.where(control_list, data_list)

ValueError: either both or neither of x and y should be given

My question is, can I do this somehow using numpy more efficiently?

Upvotes: 0

Views: 853

Answers (3)

user2357112
user2357112

Reputation: 281653

Well, first, these should be arrays:

control = np.array([1, 0, 1, 1, 0])
data = np.array([5, 4, 5, 5, 4])

Now, we can convert control to a boolean mask:

control.astype(bool)

Select the corresponding elements of data with advanced indexing:

data[control.astype(bool)]

And multiply those elements together with np.prod:

product = np.prod(data[control.astype(bool)])

Upvotes: 2

Joran Beasley
Joran Beasley

Reputation: 114038

control_list = numpy.array([1, 0, 1, 1, 0])
data_list = numpy.array([5, 4, 5, 5, 4])
numpy.product(data_list[control_list==1])

should do it ... it just says do product on datalist anywhere where control_list == 1

Upvotes: 3

Garrett R
Garrett R

Reputation: 2662

Try this out. You can convert control_list to a list of booleans, and then use it to index into data_list. Then, you can use numpy's product function to get the product of all of the values.

>>> import numpy as np
>>> cList = np.array(control_list, dtype=np.bool)
>>> cList
array([ True, False,  True,  True, False], dtype=bool)
>>> data_list = np.array(data_list)
>>> data_list[cList] # numpy supports fancy indexing
array([5, 5, 5])
>>> np.product(data_list[cList])
125

Upvotes: 3

Related Questions