Factorize values across dataframe columns with consistent mappings

Question

How can I use pandas factorize with values that exist across two columns?

Specifically, I am trying to convert values that exist in two columns to numeric values, and put the corresponding factorized values into new columns, such that the factorization is consistent with the two input columns 'A' and 'B'.

Existing DataFrame:

     A   B
0    a   b
1    c   a
2    d   a
3    e   c
4    c   b
5    b   e
6    e   f

Desired Output:

     A   B   A_ID  B_ID 
0    a   b     0     4
1    c   a     1     0
2    d   a     2     0
3    e   c     3     1
4    c   b     1     4
5    b   e     4     3
6    e   f     3     5

I am able to use factorize successfully for one column using:

df['A_ID'] = pd.factorize(df.A)[0]

How could I accomplish this with consistent mappings for values across two columns? Do I need to resort to using a custom lambda function instead, or is there a way to accomplish this with factorize?

Answer 1

pd.factorize, apply + pd.Categorical:

_, b = pd.factorize(df.values.T.reshape(-1, ))  
                           # or df.values.ravel('F'), as suggested by Zero
r = df.apply(lambda x: pd.Categorical(x, b).codes).add_suffix('_ID')

   A_ID  B_ID
0     0     4
1     1     0
2     2     0
3     3     1
4     1     4
5     4     3
6     3     5

pd.concat([df, r], 1)

   A  B  A_ID  B_ID
0  a  b     0     4
1  c  a     1     0
2  d  a     2     0
3  e  c     3     1
4  c  b     1     4
5  b  e     4     3
6  e  f     3     5

Answer 2

Let's use apply, add_suffix with pd.factorize and assign:

f = pd.factorize(df.stack().drop_duplicates().sort_index(level=1))
s1 = pd.Series(f[0], index=f[1])
print(df.assign(**df.apply(lambda x: x.map(s1)).add_suffix('_ID')))

Output:

   A  B  A_ID  B_ID
0  a  b     0     1
1  c  a     2     0
2  d  a     3     0
3  e  c     4     2
4  c  b     2     1
5  b  e     1     4

Ouput with updated dataset:

   A  B  A_ID  B_ID
0  a  b     0     4
1  c  a     1     0
2  d  a     2     0
3  e  c     3     1
4  c  b     1     4
5  b  e     4     3
6  e  f     3     5

Answer 3

Here's one way, if you want to reuse factorize values.

In [2637]: facts = np.unique(np.unique(df[['A', 'B']]), return_index=True)

In [2638]: mapping = dict(zip(*facts))

In [2639]: df.join(df[['A', 'B']].apply(lambda x: x.map(mapping)).add_suffix('_ID'))
Out[2639]:
   A  B  A_ID  B_ID
0  a  b     0     1
1  c  a     2     0
2  d  a     3     0
3  e  c     4     2
4  c  b     2     1
5  b  e     1     4
6  e  f     4     5

---

Or, using replace

In [2640]: df.join(df[['A', 'B']].replace(mapping).add_suffix('_ID'))
Out[2640]:
   A  B  A_ID  B_ID
0  a  b     0     1
1  c  a     2     0
2  d  a     3     0
3  e  c     4     2
4  c  b     2     1
5  b  e     1     4
6  e  f     4     5

---

And, to preserve your order of values use

In [2]: mapping = dict(zip(*pd.factorize(df['A'].append(df['B']).drop_duplicates())[::-1]))

In [2]: mapping
Out[2666]: {'a': 0, 'b': 4, 'c': 1, 'd': 2, 'e': 3, 'f': 5}

In [3]: df.join(df[['A', 'B']].replace(mapping).add_suffix('_ID'))
Out[3]:
   A  B  A_ID  B_ID
0  a  b     0     4
1  c  a     1     0
2  d  a     2     0
3  e  c     3     1
4  c  b     1     4
5  b  e     4     3
6  e  f     3     5

---

Details

In [2641]: facts
Out[2641]:
(array(['a', 'b', 'c', 'd', 'e', 'f'], dtype=object),
 array([0, 1, 2, 3, 4, 5], dtype=int64))

In [2642]: mapping
Out[2642]: {'a': 0, 'b': 1, 'c': 2, 'd': 3, 'e': 4, 'f': 5}