python - Pandas: create category column based on multiple columns

Question

Which would be the most efficient way to create a category column based on other columns in the row, as quickly as possible?

input:

   col1  col2  col3  col4
0     0     0   -10     1
1     1   100     0    -1
2     0     0     0     1
3     0     0   -10     1
4     1   100     0    -1

output:

   col1  col2  col3  col4 new_col
0     0     0   -10     1       1
1     1   100     0    -1       2
2     0     0     0     1       3
3     0     0   -10     1       1
4     1   100     0    -1       2

question from:https://stackoverflow.com/questions/65864251/pandas-create-category-column-based-on-multiple-columns

Answer 1

The fastest method is probably using numpy unique (if all columns are numeric):

_, new_col = np.unique(df.to_numpy(), axis=0, return_inverse=True)
df['new_col'] = new_col

or as one-liner:

df['new_col'] = np.unique(df.to_numpy(), axis=0, return_inverse=True)[1]

   col1  col2  col3  col4  new_col
0     0     0   -10     1        0
1     1   100     0    -1        2
2     0     0     0     1        1
3     0     0   -10     1        0
4     1   100     0    -1        2

---

This is about 10 times faster (for the sample data) than groupby on all columns and using the group number ngroup as category code:

df['new_col'] = df.groupby(df.columns.to_list()).ngroup()

The advantage of this method is that it also works for mixed or non-numeric typed dataframes.