python - 枚举具有相同前缀的列

Question

假设我们有以下简化数据:

df = pd.DataFrame({'A':list('abcd'),
                   'B':list('efgh'),
                   'Data_mean':[1,2,3,4],
                   'Data_std':[5,6,7,8],
                   'Data_corr':[9,10,11,12],
                   'Text_one':['foo', 'bar', 'foobar', 'barfoo'],
                   'Text_two':['bar', 'foo', 'barfoo', 'foobar'],
                   'Text_three':['bar', 'bar', 'barbar', 'foofoo']})

   A  B  Data_mean  Data_std  Data_corr Text_one Text_two Text_three
0  a  e          1         5          9      foo      bar        bar
1  b  f          2         6         10      bar      foo        bar
2  c  g          3         7         11   foobar   barfoo     barbar
3  d  h          4         8         12   barfoo   foobar     foofoo

我想枚举具有相同前缀的列。在这种情况下，前缀是 Data, Text。所以预期输出将是:

   A  B  Data_mean1  Data_std2  Data_corr3 Text_one1 Text_two2 Text_three3
0  a  e           1          5           9       foo       bar         bar
1  b  f           2          6          10       bar       foo         bar
2  c  g           3          7          11    foobar    barfoo      barbar
3  d  h           4          8          12    barfoo    foobar      foofoo

注意枚举列。

---

尝试的解决方案#1:

def enumerate_cols(dataframe, prefix):
    cols = []
    num = 1
    for col in dataframe.columns:
        if col.startswith(prefix):
            cols.append(col + str(num))
            num += 1
        else:
            cols.append(col)

    return cols

enumerate_cols(df, 'Data')

['A',
 'B',
 'Data_mean1',
 'Data_std2',
 'Data_corr3',
 'Text_one',
 'Text_two',
 'Text_three']

尝试的解决方案 #2:

[c+str(x+1) for x, c in enumerate([col for col in df.columns if col.startswith('Data')])]

['Data_mean1', 'Data_std2', 'Data_corr3']

问题:是否有更简单的解决方案，我也查看了 df.filter(like='Data') 等。但这看起来也很远获取。

---

XY问题
请确保我没有掉入 XY problem .我想用 pd.wide_to_long ，但是 stubnames 列需要以数字作为后缀才能融合数据框。

引用自文档:

With stubnames [‘A’, ‘B’], this function expects to find one or more group of columns with format A-suffix1, A-suffix2,…, B-suffix1, B-suffix2,

pd.wide_to_long(df, stubnames=['Data', 'Text'], i=['A', 'B'], j='grp', sep='_')

这将返回一个空数据框。

Answer 1

想法是将具有相同前缀的列分组，并为它们建立一个 cumcount。

由于我们需要单独处理没有前缀的列，因此我们需要使用 GroupBy.cumcount 和 np.where 分两步完成此操作:

cols = df.columns.str.split('_').str[0].to_series()

df.columns = np.where(
    cols.groupby(level=0).transform('count') > 1, 
    cols.groupby(level=0).cumcount().add(1).astype(str).radd(df.columns), 
    cols
)

df
   A  B  Data_mean1  Data_std2  Data_corr3 Text_one1 Text_two2 Text_three3
0  a  e           1          5           9       foo       bar         bar
1  b  f           2          6          10       bar       foo         bar
2  c  g           3          7          11    foobar    barfoo      barbar
3  d  h           4          8          12    barfoo    foobar      foofoo

---

一个更简单的解决方案是将您不想添加后缀的列设置为索引。然后你可以简单地做

df.set_index(['A', 'B'], inplace=True)
df.columns = (
    df.columns.str.split('_')
      .str[0]
      .to_series()
      .groupby(level=0)
      .cumcount()
      .add(1)
      .astype(str)
      .radd(df.columns))

df
     Data_mean1  Data_std2  Data_corr3 Text_one1 Text_two2 Text_three3
A B                                                                   
a e           1          5           9       foo       bar         bar
b f           2          6          10       bar       foo         bar
c g           3          7          11    foobar    barfoo      barbar
d h           4          8          12    barfoo    foobar      foofoo