Python通过大于或等于每个子组中位数的列值找出数据帧中的记录

Question

suppose I have a dataframe which could be initiated by:

df = pd.DataFrame({'group1': ['1','2','3','4','5','6'],
                   'group2': ['c','c','d','d','d','e'],
                   'value1': [1.1,2,3,4,5,6],
                   'value2': [7.1,8,9,10,11,12]
                   })
df = df.set_index(['group1', 'group2'])

I want to subset df by the value2 column, the value of which is greater or equal to the median of each sub-group specified by the index of group2 . In this example, the row of group1 in ['2','4','5','6'] should stay in the result. Can anyone help?

Answer 1

This should work:

df['value2'] = df['value2'].groupby(level='group2').transform(lambda x: np.where(x>=np.median(x), x, np.NaN))
df = df.dropna()

What this does is it gets the value2 column, and splits it into groups by group2 . For each group, it finds the median, then replaces and value below the median with NaN . It then puts this back into the value2 column, then gets rid of all the rows with NaN values.

As an alternative, here is a slightly less clear one-liner:

df = df.groupby(level='group2').transform(lambda x: x if x.name != 'group2' else np.where(x>=np.median(x), x, np.NaN)).dropna()

This does roughly the same thing, except it runs on both columns, but doesn't do anything to the group1 column.

Note that in the second approach you could instead store to a second variable, like df2 , without altering the original df if you prefer. You could do that with the first approach, but that would require yet another line to make a copy. This version is much simpler for that case.

Answer 2

I think you need to do a groupby and comparison before setting the index:

df = pd.DataFrame({'group1': ['1','2','3','4','5','6'],
                   'group2': ['c','c','d','d','d','e'],
                   'value1': [1.1,2,3,4,5,6],
                   'value2': [7.1,8,9,10,11,12]
                   })
gb = df.groupby('group2').value2.median()
df.join(gb, on='group2', rsuffix='_median')
df_filtered = df[df.value2 >= df.join(gb, on='group2', rsuffix='_median').value2_median]
df_filtered.set_index(['group1', 'group2'], inplace=True)
>>> df_filtered 
               value1  value2
group1 group2                
2      c            2       8
4      d            4      10
5      d            5      11
6      e            6      12