Python Dataframe如何根据条件创建新的列值

Question

我的数据框在下面给出

df =

index    element    data1   data2  data3 
0          M1         10      20     30
1          M1         40      50     60
2          M2         70      80     90
3          M2         100     120     130
4          M3         140     150     160
5          M3         170     180     190
6          M4         200     210     230

element_list1 = ['M1','M4',...........,'M25']
element_list2 = ['M2','M5',...........,'M26']
element_list3 = ['M3','M6',...........,'M27']

现在，我想创建一个新列。新列中的值基于元素的名称。如果元素属于list1，则选择data2，再选择list2，然后选择data2，依此类推。最后，我想实现以下类似的功能

df =

index    element    data1   data2  data3        final
0          M1         10      20     30          10
1          M1         40      50     60          40
2          M2         70      80     90          80
3          M2         100     120     130        120
4          M3         140     150     160        160
5          M3         170     180     190        190
6          M4         200     210     230        200

我的当前代码如下：

df['final'] = np.nan

for a in element_list1:
    for i,j in enumerate(df['element']):
        if j==a:
            df['final'].iloc[i] = df['data1'].iloc[i]
for a in element_list2:
    for i,j in enumerate(df['element']):
        if j==a:
            df['final'].iloc[i] = df['data2'].iloc[i]
for a in element_list3:
    for i,j in enumerate(df['element']):
        if j==a:
            df['final'].iloc[i] = df['data3'].iloc[i]

有没有一种比上面简单的方法？

Answer 1

解决方案

# Make element lists
e1 = np.arange(1,26,3)
e2 = e1 + 1
e3 = e1 + 2
element_list1 = [f'M{x}' for x in e1.tolist()]
element_list2 = [f'M{x}' for x in e2.tolist()]
element_list3 = [f'M{x}' for x in e3.tolist()]

element_lists = [element_list1, element_list2, element_list3]

# drop column 'index' from the dataframe
df = df.drop(columns='index')
# process data for 'final' column
final = pd.concat([df.data1[df.element.isin(element_lists[0])], 
                   df.data2[df.element.isin(element_lists[1])],
                   df.data3[df.element.isin(element_lists[2])],])
df['final'] = final
df

输出

制作数据

import sys
if sys.version_info[0] < 3: 
    from StringIO import StringIO
else:
    from io import StringIO

import pandas as pd

df_string = """
index    element    data1   data2  data3 
0          M1         10      20     30
1          M1         40      50     60
2          M2         70      80     90
3          M2         100     120     130
4          M3         140     150     160
5          M3         170     180     190
6          M4         200     210     230
"""

df = pd.read_csv(StringIO(df_string), sep="\s+")