In my real case I have a set of time series related to different IDs stored in a single DataFrame Some are composed by 400 samples, some by 1000 samples, some by 2000. They are stored in the same df and:

I would like to drop all the IDs made up of time series shorter than a custom length.

I wrote the following code, but I think is very ugly and inefficient.

import pandas as pd
import numpy as np



for i in set(df.id):

  if len_id>3: #3 is just a random choice for this example


   samples id
2        3  c
6        7  c
7        8  c
8        9  c
1        2  b
3        4  b
4        5  b
5        6  b

How to improve it in a more Pythonic way? Thanks


Good answer by Juho. Another option is a groupby-filter:

df.groupby('id').filter(lambda group: len(group) > 3)

#    samples id
# 1        2  b
# 2        3  c
# 3        4  b
# 4        5  b
# 5        6  b
# 6        7  c
# 7        8  c
# 8        9  c

To match your output order exactly, add a descending id sort: .sort_values('id', ascending=False)

  • 1
    \$\begingroup\$ This is neat as well. \$\endgroup\$
    – Juho
    Apr 2 at 6:19

There are many solutions. For example, you can use a groupby-transform and drop the "small" samples. An appropriate solution will depend on what your requirements are more closely, i.e., will you do the preprocessing once, and then drop different samples?

Anyway, consider:

import pandas as pd

df = pd.DataFrame({"samples": [1,2,3,4,5,6,7,8,9], "id": ["a","b","c","b","b","b","c","c","c"]})

df["counts"] = df.groupby("id")["samples"].transform("count")
df[df["counts"] > 3]

# Or if you want:
df[df["counts"] > 3].drop(columns="counts")

By the way, avoid using dict as a variable name.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.