CODEWITHSUNDEEP
pythonpandasbroadcasting
triphook
triphook

Reputation: 3097

Specifying an index for Pandas broadcasting operation

I have a Pandas Dataframe that I want to apply a condition to, and use the resulting index on another Dataframe. As a simple example:

a = pd.DataFrame(np.array([[1, 1, 2], [1, 2, 2], [2, 2, 1]]), columns=['a', 'b', 'c'])

   a  b  c
0  1  1  2
1  1  2  2
2  2  2  1

b = pd.Series([111, 222, 333], index=['a', 'b', 'c'])

a    111
b    222
c    333

In this example, I'd like to create a third array with the same dimensions as 'a', with a value of 1 (or True) where a == 2:

     a    b    c
0    0    0  333
1    0  222  333
2  111  222    0

My question is, how can I perform an operation like this and specify an axis? I suspect the solution is somehow to use pd.where, although I don't want to retain any of the original values.

Upvotes: 0

Views: 46

Answers (1)

BENY
BENY

Reputation: 323306

IIUC

a.eq(2).mul(b,axis=1)
Out[803]: 
     a    b    c
0    0    0  333
1    0  222  333
2  111  222    0

As you mentioned pd.where 

a.mask(a==1,0).mask(a==2,b,axis=1)
Out[834]: 
     a    b    c
0    0    0  333
1    0  222  333
2  111  222    0

Upvotes: 2

Related Questions