How to perform int operations in pandas dataframe with NaN elements

user_14608529 0
Viewed 50

I have a dataframe that looks something like:

    Brand    Price
0   toronto  16-Aug-16
1   quebec   18-May-17
2   brampton 18-May-17
3   toronto  31-Dec-97

My code at the moment is as follows:

df["YEAR_TORONTO"] = np.nan
df["YEAR_TORONTO"] = np.where(df["Brand"] == "toronto", df["Price"], np.nan)
df["YEAR_TORONTO"] = df["YEAR_TORONTO"].str[-2:]
df["YEAR_TORONTO"] = np.where(isinstance(df["YEAR_TORONTO"], str) and df["YEAR_TORONTO"].astype(int) >= 20, "19" + df["YEAR_TORONTO"], "20" + df["YEAR_TORONTO"])
df["YEAR_TORONTO"] = df["YEAR_TORONTO"].fillna(0).astype(int).astype(object).where(df["YEAR_TORONTO"].notnull())

This gives me a result that looks like:

    Brand    Price       YEAR_TORONTO
0   toronto  16-Aug-16   2016
1   quebec   18-May-17   NaN
2   brampton 18-May-17   NaN
3   toronto  31-Dec-97   2097

I'm aware that the 2097 instead of the 1997 is due to the and condition in my logic, but I'm not sure how to fix this due to the NaN values interfering and giving errors. Any help correcting this code would be greatly beneficial, thank you.

pandas dataframe nan numpy python
1 Answers

IIUC, you could use to_datetime and use the accessor dt.year to extract the year. Then use where to replace the rows that are not toronto with pd.NA to keep the year as integer

df["YEAR_TORONTO"] = (pd.to_datetime(df['Price']).dt.year
                        .where(df['Brand'].eq('toronto'), pd.NA))
print(df)
      Brand      Price YEAR_TORONTO
0   toronto  16-Aug-16         2016
1    quebec  18-May-17         <NA>
2  brampton  18-May-17         <NA>
3   toronto  31-Dec-97         1997
user_92747320
Related