I have this starting data set:
word1 word6 0.12 0.2
word1 word34 0.52 0.4
word1 word4 0.33 0.2
word1 word7 0.44 0.1
word2 word8 0.25 0.6
word2 word9 0.001 0.09
word2 word10 0.002 0.08
word2 word11 0.13 0.15
word3 word6 0.015 0.25
word3 word4 0.26 0.4
word3 word7 0.62 0.4
word3 word8 0.61 0.1
I want to group by column 0, and then return the rows by taking the top 2 per group, according to the third column. So here, the output should be:
word1 word6 0.12 0.2
word1 word4 0.33 0.2
word2 word9 0.001 0.09
word2 word10 0.002 0.08
word3 word6 0.015 0.25
word3 word4 0.26 0.4
In the case of equal scores, any one can be taken but that won't be an issue in this data set because in reality the numbers are given rounded to 7 decimal places.
I wrote this:
import sys
import pandas as pd
df = pd.read_csv('test_file.txt',sep='\t',names=['query','hit','normalized_value','unnormalized_value'])
df = df.sort_values(['query','normalized_value']).groupby('normalized_value').head(2)
print(df)
And the output is:
query hit normalized_value unnormalized_value
0 word1 word6 0.120 0.20
2 word1 word4 0.330 0.20
3 word1 word7 0.440 0.10
1 word1 word34 0.520 0.40
5 word2 word9 0.001 0.09
6 word2 word10 0.002 0.08
7 word2 word11 0.130 0.15
4 word2 word8 0.250 0.60
8 word3 word6 0.015 0.25
9 word3 word4 0.260 0.40
11 word3 word8 0.610 0.10
10 word3 word7 0.620 0.40
...from what I can see it's just ordered the rows by the correct column, but it hasn't returned only the top two rows per group, which i said to do with .head(2)?
Can someone explain where I'm going wrong?