如何使用正则表达式对DataFrame列标题进行分组？

Problem #

I have a DataFrame with column names structured like S1,0, S1,0.1, S1,1, etc.:

S1,0 S1,0.1 S1,0.2 S1,1 S1,1.1 S1,1.2 S2,0 S2,0.1 S2,1 S2,1.1
0    4      0      3      3      3      1      3      2      4
0    1      0      4      2      1      0      1      1      0
1    4      2      3      0      3      0      2      3      0

I want to group columns by their core prefixes: all S1,0* columns in one group, S1,1* in another, S2,0* and S2,1* similarly. Then compute aggregated stats like mean (m) and standard deviation (s), resulting in this output:

S1,0      S1,1      S2,0      S2,1
m  2.000000  2.333333  0.500000  2.500000
1  2.000000  0.666667  0.500000  2.500000
2  2.000000  1.666667  0.500000  3.000000
s  2.081666  1.154701  0.707107  2.828427
1  2.000000  0.577350  0.707107  2.121320
2  1.732051  1.527525  0.707107  0.000000

I can get this working using string splitting on column names:

import pandas as pd
import numpy as np

np.random.seed(0)
data = np.random.randint(0, 5, 30).reshape(3, 10)
df = pd.DataFrame(data, columns=['S1,0', 'S1,0.1', 'S1,0.2', 'S1,1', 'S1,1.1', 'S1,1.2', 'S2,0', 'S2,0.1', 'S2,1', 'S2,1.1'])
df = df.T
gdf = df.groupby(lambda x: x.split('.', 1)[0])[df.columns].agg({'m': np.mean, 's': np.std}).T.sort_index()

But I want to avoid split() and use regex directly. I tried passing a compiled regex to groupby(), but it throws an error:

import re
reg = re.compile('^S\d,\d')
gdf2 = df.groupby(reg)[df.columns].agg({'m': np.mean, 's': np.std}).T.sort_index()

Is there a working way to use regex for this grouping?

Solution #

Great question! You can't pass a compiled regex directly to groupby(), but you can use a lambda function that leverages regex matching to extract the grouping key from each column name. Here are two clean, regex-based approaches:

Option 1: Use re.match() in a lambda #

Define a regex pattern to capture the exact prefix you want (e.g., S1,0 or S2,1), then use re.match() to pull that match as the group key:

import pandas as pd
import numpy as np
import re

np.random.seed(0)
data = np.random.randint(0, 5, 30).reshape(3, 10)
df = pd.DataFrame(data, columns=['S1,0', 'S1,0.1', 'S1,0.2', 'S1,1', 'S1,1.1', 'S1,1.2', 'S2,0', 'S2,0.1', 'S2,1', 'S2,1.1'])
df = df.T

# Regex pattern to capture the core prefix (starts with S, followed by digit, comma, digit)
pattern = re.compile(r'^S\d,\d')

# Lambda function extracts the matched prefix for grouping
gdf = df.groupby(lambda col: pattern.match(col).group())[df.columns].agg({'m': np.mean, 's': np.std}).T.sort_index()

print(gdf)

Option 2: Use pandas str.extract() #

If you prefer using pandas' built-in string methods, you can extract the grouping keys directly from the DataFrame's index, then pass those keys to groupby():

import pandas as pd
import numpy as np

np.random.seed(0)
data = np.random.randint(0, 5, 30).reshape(3, 10)
df = pd.DataFrame(data, columns=['S1,0', 'S1,0.1', 'S1,0.2', 'S1,1', 'S1,1.1', 'S1,1.2', 'S2,0', 'S2,0.1', 'S2,1', 'S2,1.1'])
df = df.T

# Extract the core prefix using regex via str.extract()
group_keys = df.index.str.extract(r'^(S\d,\d)', expand=False)

# Group using the extracted keys
gdf = df.groupby(group_keys)[df.columns].agg({'m': np.mean, 's': np.std}).T.sort_index()

print(gdf)

How it works #

• groupby() accepts a function that takes each index label (your column name, since you transposed the DataFrame) and returns the group key. The regex-based lambda does exactly this by matching the prefix you care about.
• str.extract() is a pandas-specific method that simplifies pulling regex matches from an Index or Series. The ^(S\d,\d) pattern captures the core prefix, and expand=False returns a Series of matches that works perfectly with groupby().

Both methods will produce the exact aggregated output you want, no string splitting required.

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