Pandas finding average in a comma separated column

user_7683771 0
Viewed 207

I want to take average based on one column which is comma separated and take mean on other column.

My file looks like this:

ColumnA ColumnB
A, B, C 2.9
A, C    9.087
D       6.78
B, D, C 5.49

My output should look like this:

A   7.4435
B   5.645
C   5.83
D   6.135

My code is this:

df = pd.DataFrame(data.ColumnA.str.split(',', expand=True).stack(), columns= ['ColumnA'])  
df = df.reset_index(drop = True)
df_avg = pd.DataFrame(df.groupby(by = ['ColumnA'])['ColumnB'].mean())
df_avg = df_avg.reset_index()

It has to be around the same lines but can't figure it out.

pandas csv average
1 Answers

In your solution is created index by column ColumnB for avoid lost column values after stack and Series.reset_index, last is added as_index=False for column after aggregation:

df = (df.set_index('ColumnB')['ColumnA']
        .str.split(',', expand=True)
        .stack()
        .reset_index(name='ColumnA')
        .groupby('ColumnA', as_index=False)['ColumnB']
        .mean())
print (df)
  ColumnA   ColumnB
0       A  5.993500
1       B  4.195000
2       C  5.825667
3       D  6.135000

Or alternative solution with DataFrame.explode:

df = (df.assign(ColumnA = df['ColumnA'].str.split(','))
        .explode('ColumnA')
        .groupby('ColumnA', as_index=False)['ColumnB']
        .mean())
print (df)
  ColumnA   ColumnB
0       A  5.993500
1       B  4.195000
2       C  5.825667
3       D  6.135000
user_29010020
Related