Create a column in pandas dataframes based on conditionals

user_8340881 0
Viewed 64

I have a pandas dataframe as below:

import pandas as pd 
import numpy as np
import datetime

# intialise data of lists. 
data = {'month'      :[2,3,4,5,6,7,2,3,6,5],
        'flag': ["A","A","A","A","A","A","B","B","B","B"],
        'month1'     :[4,4,7,15,11,13,6,5,6,5],
       'value'     :[100,20,50,10,65,86,24,12,1000,200]
       } 

# Create DataFrame 
df = pd.DataFrame(data) 

# Print the output. 
df 
    month   flag    month1  value
0   2       A       4       100
1   3       A       4       20
2   4       A       7       50
3   5       A       15      10
4   6       A       11      65
5   7       A       13      86
6   2       B       6       24
7   3       B       5       12
8   6       B       6       1000
9   5       B       5       200

Now for each month in unique flag, I want to perform below logic

1) Create a variable "final" and set it to 0

2) for each month, If month1 <= max(month), set "final" for where month == month1 to "final" from month1 + value from original month. For example,

  • index 0 to 5 are one group(flag = 'A')
  • MAX of month column for group A is 7
  • for row 1(month 2), month1 is 4 which is less than 7, go to month 4(row 3) update the value of "final" column to 100(0(current "final" value)+100(value from original month)
  • perform above step to each row in a group.

Expected output:

    month   flag    month1  value   Final
0   2       A       4       100     0
1   3       A       4       20      0
2   4       A       7       50      120
3   5       A       15      10      0
4   6       A       11      65      0
5   7       A       13      86      50
6   2       B       6       24      0
7   3       B       5       12      0
8   6       B       6       1000    1024
9   5       B       5       200     212
pandas python-3.x
2 Answers

Define the following functions:

  1. A function to be applied to each row (in the current group):

    def fn(row, tbl, maxMonth):
    return tbl[tbl.month1 == row.month].value.sum()

  2. A function to be applied to each group:

    def fnGrp(grp):
    return grp.apply(fn, axis=1, tbl=grp, maxMonth=grp.month.max())

Then, to compute final column, group df by flag and apply fnGrp to each group and save the result in final column:

df['final'] = df.groupby('flag').apply(fnGrp).reset_index(level=0, drop=True)

The result (df with added column) is:

   month flag  month1  value  final
0      2    A       4    100      0
1      3    A       4     20      0
2      4    A       7     50    120
3      5    A      15     10      0
4      6    A      11     65      0
5      7    A      13     86     50
6      2    B       6     24      0
7      3    B       5     12      0
8      6    B       6   1000   1024
9      5    B       5    200    212
user_73884770

you can groupby 'flag' and 'month1' and get the sum of 'value', then merge this with df plus fillna with 0 such as:

new_df = df.merge(df.groupby(['flag', 'month1'])[['value']].sum(), 
                  left_on=['flag','month'], right_index=True, 
                  how='left', suffixes=('','_final'))\
           .fillna({'value_final':0})
print (new_df)
   month flag  month1  value  value_final
0      2    A       4    100          0.0
1      3    A       4     20          0.0
2      4    A       7     50        120.0
3      5    A      15     10          0.0
4      6    A      11     65          0.0
5      7    A      13     86         50.0
6      2    B       6     24          0.0
7      3    B       5     12          0.0
8      6    B       6   1000       1024.0
9      5    B       5    200        212.0
user_92747320
Related