Python Dataframe how to create a new column values based on a condition

Question

My dataframe is given below

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']

Now I want to create a new column. Value in the new column is based on the name of the element. If element belongs to list1 then select data2, list2 then data2, etc. Finally I want to achieve something like below

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

My present code is given below:

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]

Is there a simple approach than above?

Answer 1

Solution

# 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

Output

Make Data

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+")