Return column names as a list in dataframe for a given condition of values in python

user_5756353 0
Viewed 560

given 1xN dataframe table, need to pick 5 largest values from the row and return the corresponding column names into a list. this is the dataframe sample:

            5        2         13         15         37        8         89    
PageRank  0.444384  0.44453  0.444695  0.444882  0.444759  0.44488  0.444648

Have tried this,

r = list(pr.loc['PageRank'].nlargest(5))

but the list created has only the values in rows, not the column names. How to get the column names of 5 largest cell values? for example, in the given dataframe, it should return

[15,37,13,89,5]
pandas python-2.7 dataframe list
2 Answers

You can get some added performance by using Numpy's np.argpartition. I'll use it on the negative of the values in order to get the correct direction.

I wanted to use np.argpartition instead of sorting because it is O(n) rather than sorting which is O(nlogn).

cols = pr.columns.values
rnks = -pr.values[0]
cols[np.argpartition(rnks, 5)[:5]].tolist()

['37', '15', '13', '8', '89']

Timing
You'll notice that pir1 outperforms. But also notice that nlargest asymptotically approaches the performance of pir1 because they are both O(n).

jez1 = lambda d: list(d.loc['PageRank'].nlargest(5).index)
jez2 = lambda d: d.columns[d.loc['PageRank'].values.argsort()[::-1]][:5].tolist()
jez3 = lambda d: d.columns[d.loc['PageRank'].values.argsort()[-1:-6:-1]].tolist()
pir1 = lambda d: d.columns.values[np.argpartition(-d.values[0], 5)[:5]].tolist()

res = pd.DataFrame(
    index=[10, 30, 100, 300, 1000, 3000, 10000, 30000, 100000, 300000, 1000000],
    columns='jez1 jez2 jez3 pir1'.split(),
    dtype=float
)

for i in res.index:
    d = pd.DataFrame(dict(PageRank=np.random.rand(i))).T
    for j in res.columns:
        stmt = '{}(d)'.format(j)
        setp = 'from __main__ import d, {}'.format(j)
        res.at[i, j] = timeit(stmt, setp, number=200)

res.plot(loglog=True)

enter image description here

Timing Ratio
This table shows the ratio of each method's time relative to the minimum time taken for that particular length of array.

res.div(res.min(1), 0)

              jez1       jez2       jez3  pir1
10       20.740497   8.666576   6.738210   1.0
30       39.325125  11.962184  10.987012   1.0
100      30.121521  10.184435  10.173252   1.0
300      58.544734  11.963354  12.563072   1.0
1000     63.643729   9.361290   8.547374   1.0
3000     22.041026  15.977949  18.803516   1.0
10000     9.254778  11.620570  11.681464   1.0
30000     2.838243   7.522210   7.120721   1.0
100000    1.814005   7.486602   6.995017   1.0
300000    1.920776  13.213261  12.423890   1.0
1000000   1.332265   7.872120   7.225150   1.0
user_23366540

Use index:

r1 = list(pr.loc['PageRank'].nlargest(5).index)
print (r1)
[15, 8, 37, 13, 89]

Or:

r1 = pr.columns[pr.loc['PageRank'].values.argsort()][-1:-6:-1].tolist()
print (r1)
[15, 8, 37, 13, 89]
user_29010020
Related