Question about conditional calculation in pandas

Question

I have this formula, I wanted to turn this into pandas calculation, the formula is very easy: NEW = A(where v=1) + A(where v=3) + A(where v=5)

I have a data frame like this:

Type subType value   A           NEW
 X    a       1      3         =3+9+9=21
 X    a       3      9  
 X    a       5      9
 X    b       1      4         =4+5+0=9
 X    b       3      5 
 X    b       5      0
 Y    a       1      1         =1+2+3=6
 Y    a       3      2  
 Y    a       5      3
 Y    b       1      4         =4+5+2=11
 Y    b       3      5 
 Y    b       5      2

Two questions:

1. I know I can just write down the calculation with the specified cell, but I want the code looks nicer, is there other ways to get the value?

2. Because there will be only two results for X & Y, how can I add them into my original dataframe for further calculation? (my thought is not to add them in the dataframe and just use the value whenever it's necessary for future calculation) Quite new to coding, any answer will be appreciated!

Answer 1

Try this:

>>> import pandas as pd
>>> df = pd.DataFrame({'Type':['X','X','X','Y','Y','Y'], 'value':[1,3,5,1,3,5], 'A':[3,9,4,0,2,2]})

>>> df
  Type  value  A
0    X      1  3
1    X      3  9
2    X      5  4
3    Y      1  0
4    Y      3  2
5    Y      5  2

>>> df.groupby('Type')['A'].sum()
Type
X    16
Y     4

>>> ur_dict = df.groupby('Type')['A'].sum().to_dict()
>>> df['NEW'] = df['Type'].map(ur_dict)
>>> df
  Type  value  A  NEW
0    X      1  3   16
1    X      3  9   16
2    X      5  4   16
3    Y      1  0    4
4    Y      3  2    4
5    Y      5  2    4

Hope this helps.

Edit to answer additional inquiry:

You are mapping tuple keys to a series, that will give you an error. You should shift the columns you need to map your dictionary into as index before doing the mapping.

See below:

>>> import pandas as pd
>>> df = pd.DataFrame({'Type':['X','X','X','X','X','X','Y','Y','Y','Y','Y','Y'], 'subType':['a','a','a','b','b','b','a','a','a','b','b','b'],'value':[1,3,5,1,3,5,1,3,5,1,3,5],'A':[3,9,9,4,5,0,1,2,3,4,5,2]})
>>> df
   Type subType  value  A
0     X       a      1  3
1     X       a      3  9
2     X       a      5  9
3     X       b      1  4
4     X       b      3  5
5     X       b      5  0
6     Y       a      1  1
7     Y       a      3  2
8     Y       a      5  3
9     Y       b      1  4
10    Y       b      3  5
11    Y       b      5  2

>>> df.groupby(['Type', 'subType'])['A'].sum()
Type  subType
X     a          21
      b           9
Y     a           6
      b          11
Name: A, dtype: int64
>>> ur_dict = df.groupby(['Type', 'subType'])['A'].sum().to_dict()
>>> ur_dict
{('X', 'a'): 21, ('X', 'b'): 9, ('Y', 'a'): 6, ('Y', 'b'): 11}

>>> df['NEW'] = df.set_index(['Type', 'subType']).index.map(ur_dict)
>>> df
   Type subType  value  A  NEW
0     X       a      1  3   21
1     X       a      3  9   21
2     X       a      5  9   21
3     X       b      1  4    9
4     X       b      3  5    9
5     X       b      5  0    9
6     Y       a      1  1    6
7     Y       a      3  2    6
8     Y       a      5  3    6
9     Y       b      1  4   11
10    Y       b      3  5   11
11    Y       b      5  2   11

---