Pandas groupby mean and value_counts but using another column of counts

Question

I need a memory-efficient way of exploding a dataframe and then grouping by a column to get either normalized value_counts or the mean depending on dtype.

I have a dataframe similar to this:

   key  cnt X1  X2
0    1    8  a   1
1    1    3  b   0
2    1    4  a   0
3    2    2  b   1
4    2    6  a   0
5    3    3  a   1

The cnt variable indicates the count of the values in other columns. For example, in the first row you can interpret X1 has having 8 instances of 'a' (i.e., like ['a']*8 or ['a','a','a','a','a','a','a','a']).

I am able to do what I need using .apply() then .explode() then .groupby(), something like this:

df = pd.DataFrame([[1, 8, 'a', 1], [1, 3, 'b', 0], [1, 4, 'a', 0],
                   [2, 2, 'b', 1], [2, 6, 'a', 0], [3, 3, 'a', 1]],
                  columns=['key', 'cnt', 'X1', 'X2'])

df['X1'] = df.apply(lambda row: [row['X1']] * row['cnt'], axis=1)
df['X2'] = df.apply(lambda row: [row['X2']] * row['cnt'], axis=1)
df = df.explode(['X1', 'X2']).drop(columns=['cnt']).reset_index(drop=True)

vc = pd.DataFrame(df.groupby(['key'])['X1'].value_counts(normalize=True, dropna=True)).unstack()
vc.columns = [f'X1_{col}' for col in vc.columns.get_level_values(1).values]

df = pd.merge(left=vc.reset_index(),
              right=df.drop(columns=['X1']).groupby(['key']).agg('mean').astype(float).reset_index(),
              how='left')
print(df)

   key  X1_a  X1_b        X2
0    1  0.80  0.20  0.533333
1    2  0.75  0.25  0.250000
2    3  1.00   NaN  1.000000

But the data I'm working with is enormous, with many different variables that need to be aggregated like this, and most of the cnt values are >15000, which results in using too much memory and freezing my machine.

I feel like there must be a more memory-efficient way to do this. Anyone have any ideas?

Answer 1

There's a better way

Calculate the normalized sum of counts per key and X1

cnts = pd.crosstab(df['key'], df['X1'], df['cnt'], 
                   aggfunc='sum', normalize='index').add_prefix('X1_')

# X1   X1_a  X1_b
# key            
# 1    0.80  0.20
# 2    0.75  0.25
# 3    1.00  0.00

Calculated the weighted average of X2 per key

weighted_sum = df['cnt'].mul(df['X2']).groupby(df['key']).sum()
total_weight = df.groupby('key')['cnt'].sum()
average = weighted_sum / total_weight

# key
# 1    0.533333
# 2    0.250000
# 3    1.000000
# dtype: float64

Join the dataframes

result = cnts.join(average.rename('X2'))

#      X1_a  X1_b        X2
# key                      
# 1    0.80  0.20  0.533333
# 2    0.75  0.25  0.250000
# 3    1.00  0.00  1.000000