Search code examples
pandasdataframelookup

Horizontal lookup with sorted in pandas dataframe

I have created this pandas dataframe:

d = {'Char1': [-3,2,0], 'Char2': [0,1,2], 'Char3': [-1,0,-1]}
df = pd.DataFrame(data=d)
print(df)

which looks like this:

enter image description here

I need to create two additional fields:

  • Factor1
  • Factor2

This is how Factor1 and Factor2 should be populated across each record:

  • Factor1 should contain the name of the column with lowest value (again, across each record);
  • Factor2 should contain the name of the column with the second lowest value (again, across each record).

So, the resulting dataset should look like this:

enter image description here

So, let's take a look at the first record:

  • what is the lowest value? It's -3
  • what is the name of the column that that -3 correspond to? It's Char1 -> "Char1" is then assigned to Factor1
  • what is the second lowest value? It's -1
  • what is the name of the column that that -1 correspond to? It's Char3 -> "Char3" is then assigned to Factor2

And so on.

How can I do this in Python/Pandas?

Solution

  • An efficient method is to use the underlying numpy array with argsort:

    import numpy as np

df[['Factor1', 'Factor2']] = df.columns.to_numpy()[np.argsort(df.to_numpy())[:, :2]]

    output:

       Char1  Char2  Char3 Factor1 Factor2
0     -3      0     -1   Char1   Char3
1      2      1      0   Char3   Char2
2      0      2     -1   Char3   Char1
    generalization to N columns:
    import numpy as np

N = 2

order = np.argsort(df.to_numpy())[:, :N]
df[[f'Factor{i+1}' for i in range(N)]] = df.columns.to_numpy()[order]

    example for N=3:

       Char1  Char2  Char3 Factor1 Factor2 Factor3
0     -3      0     -1   Char1   Char3   Char2
1      2      1      0   Char3   Char2   Char1
2      0      2     -1   Char3   Char1   Char2