假设我有以下DataFrame：

arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]]
ix = pd.MultiIndex.from_arrays(arrays, names=("number", "color"))
df = pd.DataFrame([[10, 20, 30, 40], [50, 60, 70, 80]], columns=ix)
df

我想把MultiIndex的列扁平化成那样：

df.columns = [str(x[0])+"_"+str(x[1]) for x in df.columns]
df

但我想通过管道（pipe）来实现。

例如，假设函数 rename_columns 存在（它并不存在），它会是这样的：

df.rename_columns(lambda x: str(x[0])+"_"+str(x[1]))

换句话说，我想写出那个 rename_columns 函数。

我该如何实现？

函数 rename 不起作用：

df.rename(lambda x: str(x[0])+"_"+str(x[1]))

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[28], line 1
----> 1 df.rename(lambda x: str(x[0])+"_"+str(x[1]))

File .venv\Lib\site-packages\pandas\core\frame.py:6471, in DataFrame.rename(self, mapper, index, columns, axis, copy, inplace, level, errors)
   6352 """
   6353 Rename columns or index labels.
   6354 
   (...)   6468 4  3  6
   6469 """
   6470 self._check_copy_deprecation(copy)
-> 6471 return super()._rename(
   6472     mapper=mapper,
   6473     index=index,
   6474     columns=columns,
   6475     axis=axis,
   6476     inplace=inplace,
   6477     level=level,
   6478     errors=errors,
   6479 )

File .venv\Lib\site-packages\pandas\core\generic.py:1065, in NDFrame._rename(self, mapper, index, columns, axis, inplace, level, errors)
   1058             missing_labels = [
   1059                 label
   1060                 for index, label in enumerate(replacements)
   1061                 if indexer[index] == -1
   1062             ]
   1063             raise KeyError(f"{missing_labels} not found in axis")
-> 1065     new_index = ax._transform_index(f, level=level)
   1066     result._set_axis_nocheck(new_index, axis=axis_no, inplace=True)
   1068 if inplace:

File .venv\Lib\site-packages\pandas\core\indexes\base.py:6672, in Index._transform_index(self, func, level)
   6670     return type(self).from_arrays(values)
   6671 else:
-> 6672     items = [func(x) for x in self]
   6673     return Index(items, name=self.name, tupleize_cols=False)

Cell In[28], line 1
----> 1 df.rename(lambda x: str(x[0])+"_"+str(x[1]))

TypeError: 'int' object is not subscriptable

解决方案

一个简单的方法是使用 pipe+set_axis：

df.pipe(
    lambda df: df.set_axis(
        [str(x[0]) + '_' + str(x[1]) for x in df.columns], axis=1
    )
)

输出：

   1_red  1_blue  2_red  2_blue
0     10      20     30      40
1     50      60     70      80

对于通用的方法（与层级数量无关）：

def flatten_columns(df, sep='_'):
    cols = df.columns.to_frame().astype(str).agg(sep.join, axis=1).tolist()
    return df.set_axis(cols, axis=1)

df.pipe(flatten_columns)

甚至更通用，处理任意轴：

def flatten_levels(df, sep='_', axis=1):
    vals = (
        df.axes[axis].to_frame().astype(str).agg('_'.join, axis=1).tolist()
    )
    return df.set_axis(vals, axis=axis)

df.pipe(flatten_levels)
#            1_red  1_blue  2_red  2_blue
# lvl0 lvl1                              
# 1    A        10      20     30      40
#      B        50      60     70      80

df.pipe(flatten_levels, axis=0)
# number   1        2      
# color  red blue red blue
# 1_A     10   20  30   40
# 1_B     50   60  70   80