Pandas Python highest 2 rows of every 3 and tabling the results

Viewed 42

Suppose I have the following dataframe:

    .    Column1      Column2
    0       25           1
    1       89           2
    2       59           3

    3       78          10
    4       99          20
    5       38          30

    6       89         100
    7       57         200
    8       87         300

Im not sure if what I want to do is impossible or not. But I want to compare every three rows of column1 and then take the highest 2 out the three rows and assign the corresponding 2 Column2 values to a new column. The values in column 3 does not matter if they are joined or not. It does not matter if they are arranged or not for I know every 2 rows of column 3 belong to every 3 rows of column 1.

    .    Column1      Column2   Column3
    0       25           1        2
    1       89           2        3
    2       59           3
    3       78          10        20
    4       99          20        10
    5       38          30
    6       89         100        100
    7       57         200        300
    8       87         300

1 Answers

You can use np.arange with np.repeat to create a grouping array which groups every 3 values.

Then use GroupBy.nlargest then extract indices of those values using pd.Index.get_level_values, then assign them to Column3 pandas handles index alignment.

n_grps = len(df)/3
g = np.repeat(np.arange(n_grps), 3)

idx = df.groupby(g)['Column1'].nlargest(2).index.get_level_values(1)
vals = df.loc[idx, 'Column2']
vals
# 1      2
# 2      3
# 4     20
# 3     10
# 6    100
# 8    300
# Name: Column2, dtype: int64

df['Column3'] = vals
df
   Column1  Column2  Column3
0       25        1      NaN
1       89        2      2.0
2       59        3      3.0
3       78       10     10.0
4       99       20     20.0
5       38       30      NaN
6       89      100    100.0
7       57      200      NaN
8       87      300    300.0

To get output like you mentioned in the question you have to sort and push NaN to last then you have perform this additional step.

df['Column3'] = df.groupby(g)['Column3'].apply(lambda x:x.sort_values()).values

   Column1  Column2  Column3
0       25        1      2.0
1       89        2      3.0
2       59        3      NaN
3       78       10     10.0
4       99       20     20.0
5       38       30      NaN
6       89      100    100.0
7       57      200    300.0
8       87      300      NaN
Related