CODEWITHSUNDEEP
pythonpython-3.xpandasmultiplication
576i
576i

Reputation: 8362

How do I multiply a pandas column with a part of a multi index dataframe

I have a data frame with a multi index and one column.

Index fields are type and amount, the column is called count

I would like to add a column that multiplies amount and count

df2 = df.groupby(['type','amount']).count().copy()
# I then dropped all columns but one and renamed it to "count"

df2['total_amount'] = df2['count'].multiply(df2['amount'], axis='index')

doesn't work. I get a key error on amount.

How do I access a part of the multi index to use it in calculations?

Upvotes: 1

Views: 439

Answers (1)

jezrael
jezrael

Reputation: 862751

Use GroupBy.transform for Series with same size as original df with aggregated values, so possible multiple:

count = df.groupby(['type','amount'])['type'].transform('count')
df['total_amount'] = df['amount'].multiply(count, axis='index')
print (df)
   A  amount  C  D  E type  total_amount
0  a       4  7  1  5    a             8
1  b       5  8  3  3    a             5
2  c       4  9  5  6    a             8
3  d       5  4  7  9    b            10
4  e       5  2  1  2    b            10
5  f       4  3  0  4    b             4

Or:

df = pd.DataFrame({'A':list('abcdef'),
                   'amount':[4,5,4,5,5,4],
                   'C':[7,8,9,4,2,3],
                   'D':[1,3,5,7,1,0],
                   'E':[5,3,6,9,2,4],
                   'type':list('aaabbb')})

print (df)
   A  amount  C  D  E type
0  a       4  7  1  5    a
1  b       5  8  3  3    a
2  c       4  9  5  6    a
3  d       5  4  7  9    b
4  e       5  2  1  2    b
5  f       4  3  0  4    b

df2 = df.groupby(['type','amount'])['type'].count().to_frame('count')
df2['total_amount'] = df2['count'].mul(df2.index.get_level_values('amount'))
print (df2)
             count  total_amount
type amount                     
a    4           2             8
     5           1             5
b    4           1             4
     5           2            10

Upvotes: 1

Related Questions