Selecting csv row with the highest column value among rows with the same value of another column in python

Question

I created a topic model with 20 topics using Latent Dirichlet Allocation for 5000+ txt documents. I now have a .csv file which contains three columns: document number, topic number, and probability of topic in the document. It looks like this (for docs n°1 and n°2):

1   1   0,113
1   4   0,2
1   7   0,156
1   17  0,065
1   18  0,463
2   1   0,44
2   6   0,207
2   14  0,103
2   16  0,126
2   17  0,015
2   18  0,106

Basically, I want to know for a certain topic a list of the documents for which this topic has the highest probability.

I imagine I would have to do the following:

1) Get the highest value in column 3 (call it highest_prob ) for each identical value in column 1 (call it doc_number ).

2) For each of the doc_number obtained (there should be as many as there are documents), get the corresponding topic number in column 2 (call it topic_number ).

3) Return the list of doc_number associated with a particular topic_number I am interested in.

I am rather new to python and don't know how to proceed, either with the csv package or with pandas...

Answer 1

You can first replace , to . in column probability and convert to float by astype . Then groupby by column document_number and get index of max value of column probability with idxmax . Last get all records by loc :

import pandas as pd

df = pd.DataFrame({'document_number': {0: 1, 1: 1, 2: 1, 3: 1, 4: 1, 5: 2, 6: 2, 7: 2, 8: 2, 9: 2, 10: 2}, 
                   'probability': {0: '0,113', 1: '0,2', 2: '0,156', 3: '0,065', 4: '0,463', 5: '0,44', 6: '0,207', 7: '0,103', 8: '0,126', 9: '0,015', 10: '0,106'}, 
                   'topic_number': {0: 1, 1: 4, 2: 7, 3: 17, 4: 18, 5: 1, 6: 6, 7: 14, 8: 16, 9: 17, 10: 18}}, 
                    columns = ['document_number','topic_number','probability'])

print (df)
    document_number  topic_number probability
0                 1             1       0,113
1                 1             4         0,2
2                 1             7       0,156
3                 1            17       0,065
4                 1            18       0,463
5                 2             1        0,44
6                 2             6       0,207
7                 2            14       0,103
8                 2            16       0,126
9                 2            17       0,015
10                2            18       0,106

df['probability'] = df.probability.str.replace(',','.').astype(float)

print (df.groupby('document_number')['probability'].idxmax())
1    4
2    5
Name: probability, dtype: int64

print (df.loc[df.groupby('document_number')['probability'].idxmax()])
   document_number  topic_number  probability
4                1            18        0.463
5                2             1        0.440

And last set_index from column document_number and convert to to_dict columns topic_number :

print (df.loc[df.groupby('document_number')['probability'].idxmax()]
         .set_index('document_number')['topic_number'])

document_number
1    18
2     1
Name: topic_number, dtype: int64

print (df.loc[df.groupby('document_number')['probability'].idxmax()]
         .set_index('document_number')['topic_number'].to_dict())
{1: 18, 2: 1}

Another solution first sort_values by column probability and then groupby with aggregating first :

print (df.sort_values(by="probability", ascending=False)
         .groupby('document_number', as_index=False)
         .first())

   document_number  topic_number  probability
0                1            18        0.463
1                2             1        0.440

print (df.sort_values(by="probability", ascending=False)
         .groupby('document_number', as_index=False)
         .first().set_index('document_number')['topic_number'])

document_number
1    18
2     1
Name: topic_number, dtype: int64

print (df.sort_values(by="probability", ascending=False)
         .groupby('document_number', as_index=False)
         .first().set_index('document_number')['topic_number'].to_dict())
{1: 18, 2: 1}