Search code examples
pythonpandasdataframemultiple-conditions

Create new column based on conditions of other columns

I have a df like below

Year  IndexDate   WorkDate       ID   Name
0  2019        NaT 2018-12-12  9265299     FV
1  2019 2019-01-09 2019-01-09  9265299     OM
2  2020 2020-11-27 2020-11-27  9962241     PM
3  2020        NaT 2020-11-27  9962241  Other
4  2020        NaT 2021-01-19  9962241  Other
df.dtypes
Out[50]: 
Year                  int64
IndexDate    datetime64[ns]
WorkDate     datetime64[ns]
ID                    int64
Name                 object
dtype: object

df.to_dict()
{'Year': {0: 2018, 1: 2019, 2: 2020, 3: 2020, 4: 2021}, 'IndexDate': {0: NaT, 1: Timestamp('2019-01-09 00:00:00'), 2: Timestamp('2020-11-27 00:00:00'), 3: NaT, 4: NaT}, 'WorkDate': {0: Timestamp('2018-12-12 00:00:00'), 1: Timestamp('2019-01-09 00:00:00'), 2: Timestamp('2020-11-27 00:00:00'), 3: Timestamp('2020-11-27 00:00:00'), 4: Timestamp('2021-01-19 00:00:00')}, 'ID': {0: 9265299, 1: 9265299, 2: 9962241, 3: 9962241, 4: 9962241}, 'Name': {0: 'FV', 1: 'OM', 2: 'PM', 3: 'Other', 4: 'Other'}}

Each ID have one IndexDate. I'd like to create new Year column which will keep Year values if Name = OM or PM, if Name = FV or Other, new Year column will get year of the IndexDate instead of year of WorkDate

My expected result

Year  IndexDate   WorkDate       ID   Name
0  2019        NaT 2018-12-12  9265299     FV
1  2019 2019-01-09 2019-01-09  9265299     OM
2  2020 2020-11-27 2020-11-27  9962241     PM
3  2020        NaT 2020-11-27  9962241  Other
4  2020        NaT 2021-01-19  9962241  Other

Many thanks for any suggestions!!

Solution

  • If IndexDate eist only for rows with OM/PM in Name generate years and aggregate first non missing values per ID:

    df['Year'] = df['IndexDate'].dt.year.groupby(df['ID']).transform('first')

    For general solution add Series.where for set missing values for not matched OM/PM values:

    df['Year'] = (df['IndexDate'].dt.year.where(df['Name'].isin(['OM','PM']))
                             .groupby(df['ID']).transform('first'))

    IIUC need assign years by condition:

    df['Year'] = np.where(df['Name'].isin(['OM','PM']),
                      df['IndexDate'].dt.year, df['WorkDate'].dt.year)