CODEWITHSUNDEEP
pythonpandasnumpy
a0nic
a0nic

Reputation: 47

One Column Looks Up Range of Columns and Provides Corresponding Value - Python Pandas

I am attempting to create a new column called ['FirstYearSales'] which takes the values in the ['CohortYear'] column and looks up the corresponding column label and provide the corresponding value in the row. Does anyone know how to accomplish this?

data = [[2017, 150, 200, 300], [2018, 0, 750, 650], [2019, 0, 0, 50]] 
data = pd.DataFrame(data, columns = ['CohortYear', '2017', '2018', '2019']) 

    CohortYear  2017    2018    2019
0   2017        150     200     300
1   2018        0       750     650
2   2019        0       0       50

The desired results would look something like this:

    CohortYear  FirstYearSales  2017    2018    2019
0   2017        150             150     200     300
1   2018        750             0       750     650
2   2019        50              0       0       50

One of my failed attempts:

data['FirstYearSales'] = data.loc[list(data.columns.values)] == ['CohortYear']

Upvotes: 1

Views: 45

Answers (2)

anky
anky

Reputation: 75100

Try lookup with insert (for inserting column after 'CohortYear' column) after getting position from get_loc:

val = data.lookup(data.index,data['CohortYear'].map(str))
data.insert(data.columns.get_loc("CohortYear")+1,"FirstYearSales",val)

print(data)

   CohortYear  FirstYearSales  2017  2018  2019
0        2017             150   150   200   300
1        2018             750     0   750   650
2        2019              50     0     0    50

Lookup seems to be faster avoid apply on axis=1 as it is likely to be slow: (sample run against 30K rows):

m = pd.concat([data]*10000,ignore_index=True)
%%timeit
m.lookup(m.index,m['CohortYear'].map(str))
#23.7 ms ± 805 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
m
%%timeit
m.apply(lambda x: x[str(x.CohortYear)], axis=1)
#1.98 s ± 70.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Upvotes: 1

Hugolmn
Hugolmn

Reputation: 1560

Using pd.DataFrame.apply:

data['FirstYearSales'] = data.apply(lambda x: x[str(x.CohortYear)], axis=1)

   CohortYear  2017  2018  2019  FirstYearSales
0        2017   150   200   300             150
1        2018     0   750   650             750
2        2019     0     0    50              50

Upvotes: 2

Related Questions