一列查找列的范围并提供相应的值-Python Pandas

Question

我正试图创建一个名为"FirstYearSales"的新列，该列采用"CohortYear"列中的值，并查找相应的列标签并在行中提供相应的值。有人知道如何做到这一点吗？

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

想要的结果看起来像这样：

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

我失败的尝试之一：

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

Answer 1

使用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

Answer 2

从get_loc:获得位置后，尝试使用insert的lookup(用于在'CohortYear'列之后插入列(

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

查找似乎更快，避免axis=1上的apply，因为它可能很慢：(针对30K行运行示例(：

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)