Pandas: How to obtain top 2, middle 2 and bottom 2 rows in a each group

user_10278559 0
Viewed 476

Let's say I have a dataframe df as below. To obtain 1st 2 and last 2 in each group I have used groupby.nth

df = pd.DataFrame({'A': ['a','a','a','a','a','a','a','a','b','b','b','b','b','b','b'],
                   'B': [1, 2, 3, 4, 5,6,7,8,1, 2, 3, 4, 5,6,7]}, columns=['A', 'B'])
df.groupby('A').nth([0,1,-2,-1])

Result:

    B
A   
a   1
a   2
a   7
a   8
b   1
b   2
b   6
b   7

I'm not sure how to obtain the middle 2 rows. For example, in group 'A' there are 8 instances so my middle would be 4, 5 (n/2, n/2+1) and group 'B' my middle rows would be 3, 4 (n/2-0.5, n/2+0.5). Any guidance is appreciated.

pandas-groupby pandas
2 Answers

You can use iloc to find the n//2 -1 and n//2 indices for each group (// is floor division):

g = df.groupby('A')

g.apply(lambda x: x['B'].iloc[[len(x)//2-1, len(x)//2]])

A    
a  3     4
   4     5
b  10    3
   11    4
Name: B, dtype: int64
user_66711760

sacul's answer is nice , Here I just follow your own idea def a customize function 

def middle(x):
     if len(x) % 2 == 0:
         return x.iloc[int(len(x) / 2) - 1:int(len(x) / 2) + 1]
     else:
         return x.iloc[int((len(x) / 2 - 0.5)) - 1:int(len(x) / 2 + 0.5)]

pd.concat([middle(y) for _ , y in df.groupby('A')])
Out[25]:
    A  B
3   a  4
4   a  5
10  b  3
11  b  4
user_79645270
Related