CODEWITHSUNDEEP
pythonpandassorting
Tom B
Tom B

Reputation: 279

pandas sort lambda function

Given a dataframe a with 3 columns, A , B , C and 3 rows of numerical values. How does one sort all the rows with a comp operator using only the product of A[i]*B[i]. It seems that the pandas sort only takes columns and then a sort method.
I would like to use a comparison function like below. 

f = lambda i,j: a['A'][i]*a['B'][i] < a['A'][j]*a['B'][j]

Upvotes: 27

Views: 29106

Answers (3)

mork
mork

Reputation: 1863

Just adding on @srs super elegant answer an iloc option with some time comparisons with loc and the naive solution.

(iloc is preferred for when your your index is position-based (vs label-based for loc)

import numpy as np
import pandas as pd

N = 10000
df = pd.DataFrame({
                   'A': np.random.randint(low=1, high=N, size=N), 
                   'B': np.random.randint(low=1, high=N, size=N)
                  })

%%timeit -n 100
df['C'] = df['A'] * df['B']
df.sort_values(by='C')

naive: 100 loops, best of 3: 1.85 ms per loop 

%%timeit -n 100
df.loc[(df.A * df.B).sort_values().index]

loc: 100 loops, best of 3: 2.69 ms per loop

%%timeit -n 100
df.iloc[(df.A * df.B).sort_values().index]

iloc: 100 loops, best of 3: 2.02 ms per loop

df['C'] = df['A'] * df['B']

df1 = df.sort_values(by='C')
df2 = df.loc[(df.A * df.B).sort_values().index]
df3 = df.iloc[(df.A * df.B).sort_values().index]

print np.array_equal(df1.index, df2.index)
print np.array_equal(df2.index, df3.index)

testing results (comparing the entire index order) between all options:

True

True

Upvotes: 2

srs
srs

Reputation: 597

Another way, adding it here because this is the first result at Google:

df.loc[(df.A * df.B).sort_values().index]

This works well for me and is pretty straightforward. @Ami Tavory's answer gave strange results for me with a categorical index; not sure it's because of that though.

Upvotes: 16

Ami Tavory
Ami Tavory

Reputation: 76346

There are at least two ways:

Method 1

Say you start with

In [175]: df = pd.DataFrame({'A': [1, 2], 'B': [1, -1], 'C': [1, 1]})

You can add a column which is your sort key

In [176]: df['sort_val'] = df.A * df.B

Finally sort by it and drop it

In [190]: df.sort_values('sort_val').drop('sort_val', 1)
Out[190]: 
   A  B  C
1  2 -1  1
0  1  1  1

Method 2

Use numpy.argsort and then use .ix on the resulting indices:

In [197]: import numpy as np

In [198]: df.ix[np.argsort(df.A * df.B).values]
Out[198]: 
   A  B  C
0  1  1  1
1  2 -1  1

Upvotes: 28

Related Questions