DanZimmerman
DanZimmerman

Reputation: 1716

Python - multiplication based on row values

I have a data frame and three variables x,y,z.

x = 10, y = 20, z = 30

df = pd.DataFrame({'A':['a','b','c'],
                   'B':[6,7,8]})

Here is what I want to do:

Create a new column 'C':

If df['A'] == 'a', df['C'] = df['B']*x
If df['A'] == 'b', df['C'] = df['B']*y
If df['A'] == 'c', df['C'] = df['B']*z

Any quick ways to do this?

Upvotes: 2

Views: 755

Answers (4)

Allen Qin
Allen Qin

Reputation: 19947

Fastest solution so far.

#use apply to compare df.A with a,b,c and choose x,y or z.
df.apply(lambda r: r.B*[x,y,z][['a','b','c'].index(r.A)], axis=1)
Out[438]: 
0     60
1    140
2    240
dtype: int64

Speed Test

%timeit df.apply(lambda r: r.B*[x,y,z][['a','b','c'].index(r.A)], axis=1)
1000 loops, best of 3: 580 µs per loop

%timeit np.where(df.A == 'a', df.B*x, np.where(df.A == 'b', df.B*y, df.B*z))
1000 loops, best of 3: 932 µs per loop

%timeit df['B'] * df['A'].map(mul_map)
1000 loops, best of 3: 686 µs per loop

Upvotes: 0

Daan
Daan

Reputation: 942

In my opinion something more readable (albeit slower):

df.loc[df['A'] == 'a', 'C'] = df['B']*x
df.loc[df['A'] == 'b', 'C'] = df['B']*y
df.loc[df['A'] == 'c', 'C'] = df['B']*z

Upvotes: 2

root
root

Reputation: 33793

Create a dictionary mapping the column A value to the multiplication value, then use map on column A and multiply it against column B:

mul_map = {'a': 10, 'b': 20, 'c': 30}
df['C'] = df['B'] * df['A'].map(mul_map)

The resulting output:

   A  B    C
0  a  6   60
1  b  7  140
2  c  8  240

Upvotes: 4

Vaishali
Vaishali

Reputation: 38415

You can use np.where

df['C'] = np.where(df.A == 'a', df.B*x, np.where(df.A == 'b', df.B*y, df.B*z))

You get

    A   B   C
0   a   6   60
1   b   7   140
2   c   8   240

Upvotes: 2

Related Questions