[FIXED] Iterate through pandas df rows to calculate a new variable depending on conditions

Issue

I have a dataframe that looks like this:

dict1 = {'category': {0: 0.0, 1: 1.0, 2: 0.0, 3: 1.0, 4: 0.0},
 'Id': {0: 24108, 1: 24307, 2: 24307, 3: 24411, 4: 24411},
 'count': {0: 3, 1: 2, 2: 33, 3: 98, 4: 33},
 'weight': {0: 0.5, 1: 0.2, 2: 0.7, 3: 1.2, 4: 0.39}}

df1 = pd.DataFrame(dict1)

    category    Id  count   weight
0   0.0      24108    3     0.50
1   1.0      24307    2     0.20
2   0.0      24307   33     0.70
3   1.0      24411   98     1.20
4   0.0      24411   33     0.39

There are about 1000 such entries. Some only have 1 category tag, as is the example for Id 24108. Other Ids have two entries because there are two category tags, such as Id 24307 and 24411.

I want to generate a new column called val that follows one of 3 rules, depending on certain conditions.

If an Id only has 1 tag associated with it, such as Id 24108, then the value in the new column val should be a simple product between count and weight.

If an Id has 2 tags, such as Id 24307 or 24411, then the script should look at which one of the two entries has a larger count value first.

For Id 24307, the count for category 0 is higher than the count for category 1, therefore the val column should be 2 * (0.2+1) next to category 1 and 33 * (1-0.7) next to category 0

For Id 24411, the count for category 0 is lower than the count for category 1, therefore the val column should be 98 * (1-1.2) next to category 1 and 33 * (0.3+1) next to category 0.

The expected output should be like this:

    category    Id  count   weight  val
0   0.0      24108    3     0.50   1.50
1   1.0      24307    2     0.20   2.40
2   0.0      24307   33     0.70   9.90
3   1.0      24411   98     1.20 -19.60
4   0.0      24411   33     0.39  42.90

Solution

You can use boolean arithmetics and numpy.select:

g = df1.groupby('Id')

# number of categories
n = g['category'].transform('nunique')
# max count
m = g['count'].transform('max')

# selection based on conditions
df1['val'] = np.select(
          [n.eq(1), df1['count'].eq(m)], # case if unique, case if max
          # formula for unique, formula for max
          [df1['count']*df1['weight'], df1['count']*(1-df1['weight'])],
          # default value (i.e. not unique and non-max count)
          df1['count']*(1+df1['weight'])
         )

output:

   category     Id  count  weight    val
0       0.0  24108      3    0.50   1.50
1       1.0  24307      2    0.20   2.40
2       0.0  24307     33    0.70   9.90
3       1.0  24411     98    1.20 -19.60
4       0.0  24411     33    0.39  45.87

If you really have only 2 categories, then you can simplify the arithmetic to consider the non max as min. It's a bit quite hacky but should work fine:

g = df1.groupby('Id')

# 0 if unique category else 1
n = g['category'].transform('nunique').gt(1).astype(int)
# -1 if count is min of group else 1
m = (df1['count'] != g['count'].transform('min'))*2-1

df1['val'] = df1['count'] * (n-m*df1['weight'])

Answered By - mozway