How to impute NaN values based on values of other column?

Question

I have 2 columns in dataframe

1)work experience (years)

2)company_type

I want to impute company_type column based on work experience column. company_type column has NaN values which I want to fill based on work experience column. Work experience column does not have any missing values.

Here work_exp is numerical data and company_type is categorical data.

Example data:

Work_exp      company_type
   10            PvtLtd
   0.5           startup
   6           Public Sector
   8               NaN
   1             startup
   9              PvtLtd
   4               NaN
   3           Public Sector
   2             startup
   0               NaN 

I have decided the threshold for imputing NaN values.

Startup if work_exp < 2yrs
Public sector if work_exp > 2yrs and <8yrs
PvtLtd if work_exp >8yrs

Based on above threshold criteria how can I impute missing categorical values in column company_type.

Answer 1

You can use numpy.select with numpy.where:

# define conditions and values
conditions = [df['Work_exp'] < 2, df['Work_exp'].between(2, 8), df['Work_exp'] > 8]
values = ['Startup', 'PublicSector', 'PvtLtd']

# apply logic where company_type is null
df['company_type'] = np.where(df['company_type'].isnull(),
                              np.select(conditions, values),
                              df['company_type'])

print(df)

   Work_exp  company_type
0      10.0        PvtLtd
1       0.5       startup
2       6.0  PublicSector
3       8.0  PublicSector
4       1.0       startup
5       9.0        PvtLtd
6       4.0  PublicSector
7       3.0  PublicSector
8       2.0       startup
9       0.0       Startup

pd.Series.between includes start and end values by default, and permits comparison between float values. Use inclusive=False argument to omit boundaries.

s = pd.Series([2, 2.5, 4, 4.5, 5])

s.between(2, 4.5)

0     True
1     True
2     True
3     True
4    False
dtype: bool

Answer 2

great answer by @jpp. Just want to add a different approach here using pandas.cut().

df['company_type'] = pd.cut(
    df.Work_exp,
    bins=[0,2,8,100],
    right=False,
    labels=['Startup', 'Public', 'Private']
)



   Work_exp company_type
0   10.0    Private
1   0.5     Startup
2   6.0     Public
3   8.0     Private
4   1.0     Startup
5   9.0     Private
6   4.0     Public
7   3.0     Public
8   2.0     Public
9   0.0     Startup

Also based on your conditions, Index 8 should be public ?

• Startup < 2
• PublicSector >=2 and < 8
• PvtLtd >= 8