Panda groupby: считать строки, удовлетворяющие условию на других столбцах?

Question

Я хотел бы сделать groupby в пандах, получая в результате фрейм данных, который имеет в качестве столбцов столбец, используемый для groupby, количество элементов для каждой группы и среди них количество элементов в каждой группе, количество элементов, которое удовлетворяет / не удовлетворяет условию на основе значения другого столбца.

Например, входные данные, подобные этому:

type    success
A       True
B       False
A       False
C       True

Я бы хотел что-то вроде:

type    total    numOfSuccess numOfFailure
A       2        1             1
B       1        0             1
C       1        1             0

В pyspark я сделал это как

import pyspark.sql.functions as F
df = df.groupBy("type").agg(\
    F.count('*').alias('total'), \
    F.sum(F.when(F.col('success')=="true", 1).otherwise(0)).alias('numOfSuccess'),
    F.sum(F.when(F.col('success')!="true", 1).otherwise(0)).alias('numOfFails'))

пока в пандах я могу получить только total и numOfSuccess как:

df_new = df.groupby(['type'], as_index=False)['success'].agg({'total':'count', 'numOfSuccess':'sum'})

или только сумма как:

df = df.groupby(['type']).size().reset_index(name='NumOfReqs')

но я не могу получить третий столбец numOfFailures и плюс, если есть альтернатива, а не суммирование логических значений, было бы лучше, поскольку, по моему мнению, она может быть распространена и на другие случаи.

Как я могу это сделать?

Answer 1

Используйте groupby с GroupBy.size для подсчета всех данных, затем для подсчета по категориям нужно поворачиваться - с GroupBy.size и unstack , crosstab или pivot_table:

df1 = df.groupby('type').size().reset_index(name='count')
df2 = (df.groupby(['type', 'success']).size().unstack(fill_value=0)
        .rename(columns={True:'numOfSuccess', False:'numOfFails'}))

Альтернатива для df2:

df2 = pd.crosstab(df['type'], df['success'])
        .rename(columns={True:'numOfSuccess', False:'numOfFails'}))

Или:

df2 = (df.pivot_table(index='type', columns='success', fill_value=0, aggfunc='size')
        .rename(columns={True:'numOfSuccess', False:'numOfFails'}))

---

df_new = df1.join(df2, on='type')
print (df_new)
  type  count  numOfFails  numOfSuccess
0    A      2           1             1
1    B      1           1             0
2    C      1           0             1

Еще одним решением является использование параметра margins в crosstab и удаление последней строки путем индексации с помощью iloc:

df = (pd.crosstab(df['type'], df['success'], margins=True)
        .rename(columns={True:'numOfSuccess', False:'numOfFails', 'All':'count'})
        .iloc[:-1]
        .reset_index()
        .rename_axis(None, axis=1))

print (df)
  type  numOfFails  numOfSuccess  count
0    A           1             1      2
1    B           1             0      1
2    C           0             1      1

РЕДАКТИРОВАТЬ: Если возможно True или False не существует, добавьте reindex для добавления отсутствующего столбца:

print (df)
  type  success
0    A     True
1    B     True
2    A     True
3    C     True

df1 = df.groupby('type').size().reset_index(name='count')
df2 = (df.groupby(['type', 'success']).size().unstack(fill_value=0)
         .reindex(columns=[True, False], fill_value=0)
         .rename(columns={True:'numOfSuccess', False:'numOfFails'}))


df_new = df1.join(df2, on='type')
print (df_new)
  type  count  numOfSuccess  numOfFails
0    A      2             2           0
1    B      1             1           0
2    C      1             1           0