Pandas DataFrame: Simple assignment - the floor of the difference of two column values to another column value, loop does not work

Question

How would you perform the following in Pandas?

import math
for index, row in data.iterrows():
  if row["year"] == 0:
    row["year"] = math.floor((row["death"] - row["birth"])/2)

This loop does not work but what i am trying to do is to assign the floor of the difference divided by 2 of the death and birth columns to the column year if the column year has the value 0. I know you should avoid looping in Pandas and this probably has a simple solution but i can not figure it out right now.

Answer 1

Use numpy.where :

import numpy as np

df['newcol'] = np.where(df['year'] == 0, math.floor((df['death'] - df['birth'])/2), df['year'])

This essentially is:

np.where(condition, if True then, if False then)

Answer 2

Slice with loc :

df.loc[df['year'] == 0, 'year'] = np.floor((df.loc[df['year'] == 0, 'death'] - df.loc[df['year'] == 0, 'birth']) / 2)

Maybe a more readable solution:

mask = df['year'] == 0
df.loc[mask, 'year'] = np.floor((df.loc[mask, 'death'] - df.loc[mask, 'birth']) / 2)