Dance Party2
Dance Party2

Reputation: 7536

Pandas assign value of one column based on another

Given the following data frame:

import pandas as pd
df = pd.DataFrame(
        {'A':[10,20,30,40,50,60],
         'B':[1,2,1,4,5,4]
        })
df

    A   B
0   10  1
1   20  2
2   30  1
3   40  4
4   50  5
5   60  4

I would like a new column 'C' to have values be equal to those in 'A' where the corresponding values for 'B' are less than 3 else 0. The desired result is as follows:

    A   B  C
0   10  1  10
1   20  2  20
2   30  1  30
3   40  4  0
4   50  5  0
5   60  4  0

Thanks in advance!

Upvotes: 3

Views: 7161

Answers (2)

EdChum
EdChum

Reputation: 394061

Here you can use pandas method where direct on the column:

In [3]:
df['C'] = df['A'].where(df['B'] < 3,0)
df

Out[3]:
    A  B   C
0  10  1  10
1  20  2  20
2  30  1  30
3  40  4   0
4  50  5   0
5  60  4   0

Timings

In [4]:
%timeit df['A'].where(df['B'] < 3,0)
%timeit np.where(df['B'] < 3, df['A'], 0)

1000 loops, best of 3: 1.4 ms per loop
1000 loops, best of 3: 407 µs per loop

np.where is faster here but pandas where is doing more checking and has more options so it depends on the use case here.

Upvotes: 4

mechanical_meat
mechanical_meat

Reputation: 169334

Use np.where:

df['C'] = np.where(df['B'] < 3, df['A'], 0)

>>> df
    A  B   C
0  10  1  10
1  20  2  20
2  30  1  30
3  40  4   0
4  50  5   0
5  60  4   0

Upvotes: 6

Related Questions