Python Pandas: create variable for unique combinations of 2 categorical variables?

Viewed 1435

Say I have some data:

df = pd.DataFrame({'location':['store','online','store','online','online'],
                  'item': ['apple','apple','orange','orange','orange']})
df
>>>

location    item
0   store   apple
1   online  apple
2   store   orange
3   online  orange
4   online  orange

As you'll note there are four possible variable combinations: (store,apple), (online,apple), (store,orange), (online,orange). I'd like to assign a dummy variable column. My naive approach creates four dummy variables, whereas I want a single label column:

pd.get_dummies(df,['location','item'])
>>>

location_online location_store  item_apple  item_orange
0   0   1   1   0
1   1   0   1   0
2   0   1   0   1
3   1   0   0   1
4   1   0   0   1

Whereas I'd prefer it to look like:

df 
>>>
location    item   combination     dummy
0   store   apple  (store, apple)   0
1   online  apple  (online, apple)  1
2   store   orange (store, orange)  2
3   online  orange (online, orange) 3
4   online  orange (online, orange) 3

Note, the dummy only equals the index because there are only 4 rows. This obviously would not be universally true.

Edit1: Above edited in response to comment. Edit2: I've added a 5th row to illustrate that a row can be repeated, however, it should have the same dummy/combination as its duplicate.

3 Answers

Let's create combinations by concatenating location and item then use factorize to encode these combinations to get dummy variables:

df['combination'] = df['location'].add(', ' + df['item'])
df['dummy'] = df['combination'].factorize()[0]

  location    item     combination  dummy
0    store   apple    store, apple      0
1   online   apple   online, apple      1
2    store  orange   store, orange      2
3   online  orange  online, orange      3
4   online  orange  online, orange      3

Let's try agg:

df['combination'] = df[['location','item']].agg(tuple, axis=1)
df['dummy'] = df['combination'].factorize()[0]

Output:

  location    item       combination  dummy
0    store   apple    (store, apple)      0
1   online   apple   (online, apple)      1
2    store  orange   (store, orange)      2
3   online  orange  (online, orange)      3
4   online  orange  (online, orange)      3

If you don't care about combination, you can use groupby.ngroup():

df['dummy'] = df.groupby(['location','item'], sort=False).ngroup()

Output:

  location    item  dummy
0    store   apple      0
1   online   apple      1
2    store  orange      2
3   online  orange      3
4   online  orange      3

You can apply a lambda function on first 2 columns. See below. d is a dictionary with dummies per pair.

d={('store', 'apple'):0, ('online', 'apple'):1, ('store', 'orange'):2, ('online', 'orange'):3}

def f(x,y):
    return d[(x,y)]

df['dummy'] = df[['location','item']].apply(lambda x: f(*x), axis=1)

>>>print(df)
  location    item  dummy
0    store   apple      0
1   online   apple      1
2    store  orange      2
3   online  orange      3
4   online  orange      3
Related