6
\$\begingroup\$

I want to get exactly n unique randomly sampled rows per category in a Dataframe. This proved to be involve more steps than the description would lead you to believe.

n = 4
df = pd.DataFrame({'category': [1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 1],
                   'value' : range(12)})


category_counts = df['category'].value_counts()
categories_with_count_above_threshold = category_counts[category_counts >= n].index

# rows with infrequent categories are filtered out
df = df[df['category'].isin(categories_with_count_above_threshold)]

# sample exactly x rows per category
df = df.groupby('category').sample(n)

This goes through the whole DataFrame quite a few times. With a bigger DataFrame this can become quite time consuming. Is there a way to simplify this expression?

PS: requires pandas >= 1.1 for the DataFrameGroupBy.sample

\$\endgroup\$
0
\$\begingroup\$

Have you tried using DataFrameGroupBy.filter() (I don't have recent pandas on this machine so I can't test it).

n = 4
df = pd.DataFrame({'category': [1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 1],
                   'value' : range(12)})

# sample exactly x rows per category
df = df.groupby('category').filter(lambda x:x.size >= n).sample(n)
\$\endgroup\$
2
  • \$\begingroup\$ Filter sadly does not return groups, so sample works on the whole DataFrame and I get n rows back instead of n rows per group. As it stands, I would just have to do another groupby('category') between filter and sample. \$\endgroup\$ – Kaio Sep 14 '20 at 11:28
  • \$\begingroup\$ Also, filter passes the groups as DFs to the lambda function and size returns <n_rows * n_columns> for DFs, so you would have to call len instead. In the end the multiple groupby could not be avoided and the query has to look like this afaict: df = df.groupby('category').filter(lambda x: len(x) >= n).groupby('category').sample(n) \$\endgroup\$ – Kaio Sep 15 '20 at 7:47
0
\$\begingroup\$

Here's code using .apply(), that makes a single pass through the dataframe.

df = pd.DataFrame({'category': [1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 1],
                   'value': range(12),
                   'other': list('abcdefghijkl')})

def pick(group, n):
    """Return n random rows from groups that have at least n rows"""

    if len(group) >= n:
        return group.sample(n)

# sample exactly x rows per category, iff there are at least x in the catgory
x = 4
df1 = df.groupby('category').apply(pick, x).reset_index(drop=True)

df1

output:

    category    value   other
0      1         3      d
1      1        11      l
2      1         6      g
3      1         9      j
4      2        10      k
5      2         1      b
6      2         4      e
7      2         7      h
\$\endgroup\$

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.