0

I want to create a new column in python dataframe based on other column values in multiple rows. For example, my python dataframe df:

A    |    B
------------
10   |    1
20   |    1
30   |    1
10   |    1
10   |    2
15   |    3
10   |    3

I want to create variable C that is based on the value of variable A with condition from variable B in multiple rows. When the value of variable B in row i,i+1,..., the the value of C is the sum of variable A in those rows. In this case, my output data frame will be:

  A   |   B   |   C
--------------------
  10  |   1   |   70
  20  |   1   |   70
  30  |   1   |   70
  10  |   1   |   70
  10  |   2   |   10
  15  |   3   |   25
  10  |   3   |   25

I haven't got any idea the best way to achieve this. Can anyone help?

Thanks in advance

Supratim Haldar
  • 2,376
  • 3
  • 16
  • 26
Bernando Purba
  • 541
  • 2
  • 9
  • 18

2 Answers2

1

recreate the data:

import pandas as pd

A = [10,20,30,10,10,15,10]

B = [1,1,1,1,2,3,3]

df = pd.DataFrame({'A':A, 'B':B})

df

     A   B
0   10   1
1   20   1
2   30   1
3   10   1
4   10   2
5   15   3
6   10   3

and then i'll create a lookup Series from the df:

lookup = df.groupby('B')['A'].sum()
lookup

    A
B
1   70
2   10
3   25

and then i'll use that lookup on the df using apply

df.loc[:,'C'] = df.apply(lambda row: lookup[lookup.index == row['B']].values[0], axis=1)
df

         A   B    C
    0   10   1   70
    1   20   1   70
    2   30   1   70
    3   10   1   70
    4   10   2   10
    5   15   3   25
    6   10   3   25
zero
  • 1,605
  • 3
  • 15
  • 24
0

You have to use groupby() method, to group the rows on B and sum() on A.

df['C'] = df.groupby('B')['A'].transform(sum)
Supratim Haldar
  • 2,376
  • 3
  • 16
  • 26