Calculate specific sums in dataframe, based on unique values in two other columns, and write to new column

Viewed 292

I have a date frame with 3 columns, including some duplicate rows:

dict1 = {'experiment': ['A', 'B', 'B', 'C', 'C', 'C', 'C'], 'run': ['A-1', 'B-1', 'B-2', 'C-1', 'C-1', 'C-2', 'C-2'], 'data': [6, 5, 5, 4, 4, 4, 4]}
df1 = pd.DataFrame(data=dict1)
print(df1)

  experiment  run  data
0          A  A-1     6
1          B  B-1     5
2          B  B-2     5
3          C  C-1     4
4          C  C-1     4
5          C  C-2     4
6          C  C-2     4

I am trying to create a new column that, for each row, contains the sum of column 'data' of unique runs for that experiment. The duplicate rows should stay intact. So my expected outcome is:

  experiment  run  data  exp-sum
0          A  A-1     6        6
1          B  B-1     5       10
2          B  B-2     5       10
3          C  C-1     4        8
4          C  C-1     4        8
5          C  C-2     4        8
6          C  C-2     4        8

I have tried combining .groupby and .unique, but so far I only get the correct sums per run, which would need to be further summed up (per experiment) and then written into the original df.

print(df1.groupby('run')['data'].unique())

run
A-1    [6]
B-1    [5]
B-2    [5]
C-1    [4]
C-2    [4]

Any input very welcome!

3 Answers

If I get the objective right, the code below should do the job

  1. Get unique experiments
unique = df.drop_duplicates(subset=['experiment', 'run'], keep='first')
  1. Get sum of data per unique experiment
sums = pd.DataFrame( unique.groupby('experiment').data.sum() ).reset_index(drop=False)
  1. Add new computed column to the original df through a join
df = df.merge( sums, on=['experiment'], how='inner' )

Another solution, using .pivot_table:

df1 = df1.set_index("experiment")
x = df1.pivot_table(
    index=pd.Grouper(level=0),
    columns="run",
    values="data",
    aggfunc=lambda x: x.unique().sum(),
).sum(axis=1)
df1["exp-sum"] = x
print(df1.reset_index())

Prints:

  experiment  run  data  exp-sum
0          A  A-1     6      6.0
1          B  B-1     5     10.0
2          B  B-2     5     10.0
3          C  C-1     4      8.0
4          C  C-1     4      8.0
5          C  C-2     4      8.0
6          C  C-2     4      8.0

You can create a map of unique runs of experiment as key and sum of data as values. Then use Series.map to map the values to experiment,

mapper = df1.drop_duplicates('run').groupby('experiment')['data'].sum()

print(mapper)

experiment
A     6
B    10
C     8

df1['exp-sum'] = df1['experiment'].map(mapper)

print(df1)

    experiment  run   data  exp-sum
0   A           A-1   6     6
1   B           B-1   5     10
2   B           B-2   5     10
3   C           C-1   4     8
4   C           C-1   4     8
5   C           C-2   4     8
6   C           C-2   4     8
Related