如何使用apply函数重命名pandas dataframe中的列？

import pandas as pd

# 创建一个示例DataFrame
df = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})

# 定义一个函数来重命名列
def rename_column(column_name):
    if column_name == 'A':
        return 'New_A'
    elif column_name == 'B':
        return 'New_B'
    else:
        return column_name

# 使用apply函数调用重命名函数，并更新列名
df = df.rename(columns=lambda x: rename_column(x))

# 打印更新后的DataFrame
print(df)

   New_A  New_B
0      1      4
1      2      5
2      3      6