Удаление случайных строк из фрейма данных до тех пор, пока число не станет равным некоторым критериям

Question

У меня есть фрейм данных с данными, которые я передаю в библиотеку ML на python. Данные, которые я имею, подразделяются на 5 различных задач: t1, t2, t3, t4, t5. Данные, которые у меня есть сейчас для каждой задачи, являются неравномерными, для упрощения приведем пример.

task, someValue
t1,   XXX
t1,   XXX
t1,   XXX
t1,   XXX
t2,   XXX
t2,   XXX

В приведенном выше случае я хочу удалить случайные строки с меткой задачи «t1», пока не будет равно «t1», как «t2» Поэтому после запуска кода он должен выглядеть следующим образом:

task, someValue
t1,   XXX
t1,   XXX
t2,   XXX
t2,   XXX

Какой самый чистый способ сделать это? Я мог бы, конечно, просто сделать для циклов и, если условия и использовать случайные числа и подсчитать вхождения для каждой итерации, но это решение не будет очень элегантным. Наверняка должен быть способ использования функций датафрейма? Пока что вот что я получил:

def equalize_rows(df):
    t = df['task'].value_counts()
    mininmum_occurance = min(t)

Answer 1

Вы можете вычислить наименьшее количество задач в вашем фрейме данных, а затем использовать groupby + head, чтобы получить первые N строк для каждой задачи.

v = df['task'].value_counts().min()
df = df.groupby('task', as_index=False).head(v)

df
  task someValue
0   t1       XXX
1   t1       XXX
4   t2       XXX
5   t2       XXX