n回以上繰り返されたとき、Pandasデータフレームの内の連続する重複をn回まで残して、以降は削除する方法

Question

前提・実現したいこと

pythonを用いてデータ整理を行っています。
データフレーム内の指定された回数以上同じ数が連続している場合、連続n行目までを残してn行以降は削除する方法が知りたいです。

調べると、https://www.fixes.pub/program/32835.htmlの回答が一番目標に近かったのですが、ここではn行以上続いた場合、連続するまとまりの行すべてが削除されてしまっていました。

私の目標を具体的に説明すると、このようなDataFrameがある場合、

x  y
1  2
2  2
3  3
4  3
5  3
6  3
7  3
8  3
9  3
10 4
11 4
12 4
13 2

yの値が5行以上連続しているものは5行目まで残し、下記のようなデータフレームを出力したいと思っています。

x  y
1  2
2  2
3  3
4  3
5  3
6  3
7  3
10 4
11 4
12 4
13 2

試したこと

https://www.fixes.pub/program/32835.htmlでは、

thresh = 5
s = df['y'].diff().ne(0).cumsum()
small_size = s.groupby(s).transform('size') < thresh
first_rows = ~s.duplicated()
df[small_size | first_rows]

というコードが記載されており、設定する数字や演算子をいじってみたり、新たな条件（s.groupby(s).transform('size') > thresh の場合）を加えて指定された行数だけ残す式を考えようとしましたがわからず、目標の出力は行えませんでした。

python初心者で、関数に関して不明瞭な点が多く、データ整理が滞っています。
お力添えいただけましたら幸いです。よろしくお願いいたします。

Answer 1

グループ分けまではできているので、あとはGroupby.head() を使って、グループごとに先頭からthresh行まで取ってくるようにすればいいです。

thresh = 5
s = df['y'].diff().ne(0).cumsum()
print(df.groupby(s).head(thresh))

Comments

[towai]

ありがとうございます。目標のデータが出力できました。