如何高效统计pandas多分组DataFrame的子组计数并映射回原表

Problem Context #

I have a pandas DataFrame with columns "A" and "B" as follows:

import pandas as pd

data = {
    'A': ['foo', 'bar', 'foo', 'bar', 'foo', 'bar', 'foo', 'foo', 'xyz'],
    'B': ['one', 'one', 'two', 'one', 'two', 'two', 'one', 'one', 'one']
}
df = pd.DataFrame(data)

I want to calculate, for each group in column "A", the count of each value in column "B" (specifically counts for "one" and "two"), then add these counts as new columns (countOne, countTwo) to every row in the original DataFrame. The desired output looks like this:

My initial approach using groupby + map with lambdas is extremely slow, so I'm looking for a more efficient solution.

Efficient Solution #

We can achieve this with grouped aggregation + pivoting + merging—all vectorized operations that leverage pandas' optimized backend:

1. Group and pivot to get subgroup counts
   First, group by "A" and "B" to get counts, then pivot the result into a wide-format table and fill missing values with 0 (so groups without "two" get a 0 count):

   # Calculate counts for each (A, B) pair, then pivot to wide format
   group_counts = df.groupby(['A', 'B']).size().unstack(fill_value=0)
   # Rename columns to match our desired output
   group_counts = group_counts.rename(columns={'one': 'countOne', 'two': 'countTwo'})
   

   This gives us a table where each row is a group from "A", with precomputed counts:

   countOne  countTwo
   A                      
   bar        2         1
   foo        3         2
   xyz        1         0
   

2. Merge counts back to the original DataFrame
   Use merge to attach the precomputed counts to every row in the original DataFrame, matching on column "A":

   result_df = df.merge(group_counts, on='A', how='left')
   

Full Working Code #

import pandas as pd

# Create original DataFrame
data = {
    'A': ['foo', 'bar', 'foo', 'bar', 'foo', 'bar', 'foo', 'foo', 'xyz'],
    'B': ['one', 'one', 'two', 'one', 'two', 'two', 'one', 'one', 'one']
}
df = pd.DataFrame(data)

# Generate grouped counts
group_counts = df.groupby(['A', 'B']).size().unstack(fill_value=0)
group_counts = group_counts.rename(columns={'one': 'countOne', 'two': 'countTwo'})

# Merge to get final result
result_df = df.merge(group_counts, on='A', how='left')

print(result_df)

Why This Is Better #

• Avoids slow row-wise operations (like map with lambdas) which perform poorly on large datasets.
• Uses pandas' optimized C-backed operations for grouping, pivoting, and merging—these are far faster than manual Python-level loops.
• Ensures edge cases (like the "xyz" group with no "two" values) are handled cleanly with fill_value=0.

内容的提问来源于stack exchange，提问作者jessepeng