python pandas, DF.groupby().agg(), column reference in agg()

agg is the same as aggregate. It’s callable is passed the columns (Series objects) of the DataFrame, one at a time.


You could use idxmax to collect the index labels of the rows with the maximum
count:

idx = df.groupby('word')['count'].idxmax()
print(idx)

yields

word
a       2
an      3
the     1
Name: count

and then use loc to select those rows in the word and tag columns:

print(df.loc[idx, ['word', 'tag']])

yields

  word tag
2    a   T
3   an   T
1  the   S

Note that idxmax returns index labels. df.loc can be used to select rows
by label. But if the index is not unique — that is, if there are rows with duplicate index labels — then df.loc will select all rows with the labels listed in idx. So be careful that df.index.is_unique is True if you use idxmax with df.loc


Alternative, you could use apply. apply‘s callable is passed a sub-DataFrame which gives you access to all the columns:

import pandas as pd
df = pd.DataFrame({'word':'a the a an the'.split(),
                   'tag': list('SSTTT'),
                   'count': [30, 20, 60, 5, 10]})

print(df.groupby('word').apply(lambda subf: subf['tag'][subf['count'].idxmax()]))

yields

word
a       T
an      T
the     S

Using idxmax and loc is typically faster than apply, especially for large DataFrames. Using IPython’s %timeit:

N = 10000
df = pd.DataFrame({'word':'a the a an the'.split()*N,
                   'tag': list('SSTTT')*N,
                   'count': [30, 20, 60, 5, 10]*N})
def using_apply(df):
    return (df.groupby('word').apply(lambda subf: subf['tag'][subf['count'].idxmax()]))

def using_idxmax_loc(df):
    idx = df.groupby('word')['count'].idxmax()
    return df.loc[idx, ['word', 'tag']]

In [22]: %timeit using_apply(df)
100 loops, best of 3: 7.68 ms per loop

In [23]: %timeit using_idxmax_loc(df)
100 loops, best of 3: 5.43 ms per loop

If you want a dictionary mapping words to tags, then you could use set_index
and to_dict like this:

In [36]: df2 = df.loc[idx, ['word', 'tag']].set_index('word')

In [37]: df2
Out[37]: 
     tag
word    
a      T
an     T
the    S

In [38]: df2.to_dict()['tag']
Out[38]: {'a': 'T', 'an': 'T', 'the': 'S'}

Leave a Comment

Hata!: SQLSTATE[HY000] [1045] Access denied for user 'divattrend_liink'@'localhost' (using password: YES)