如何在Python的Pandas DataFrame中实现类似R的dplyr::relocate列重排？

我之前也在找Pandas里有没有类似R语言dplyr包中relocate()的简便方法，用来重排DataFrame的列，搜了一圈没找到合适的方案，于是和同事一起写了个简单的Python函数来实现这个功能，分享给需要的人：

import pandas as pd

def relocate(df, col, position='first', ref_col=None):
    """
    在Pandas DataFrame中重排指定列，模拟dplyr的relocate函数
    
    参数:
        df: pandas.DataFrame - 要操作的DataFrame
        col: str - 需要移动的列名
        position: str - 移动位置，可选值为'first', 'last', 'before', 'after'
        ref_col: str - 当position为'before'或'after'时，指定参考列名
    
    返回:
        pandas.DataFrame - 列重排后的DataFrame
    """
    if col not in df.columns:
        raise ValueError(f"列名 {col} 不存在于DataFrame中")
    
    cols = list(df.columns)
    cols.remove(col)
    
    if position == 'first':
        new_cols = [col] + cols
    elif position == 'last':
        new_cols = cols + [col]
    elif position == 'before':
        if ref_col not in cols:
            raise ValueError(f"参考列名 {ref_col} 不存在于DataFrame中")
        idx = cols.index(ref_col)
        new_cols = cols[:idx] + [col] + cols[idx:]
    elif position == 'after':
        if ref_col not in cols:
            raise ValueError(f"参考列名 {ref_col} 不存在于DataFrame中")
        idx = cols.index(ref_col)
        new_cols = cols[:idx+1] + [col] + cols[idx+1:]
    else:
        raise ValueError("position参数只能是'first', 'last', 'before', 'after'中的一个")
    
    return df[new_cols]

使用示例 #

# 示例DataFrame
df = pd.DataFrame({
    'A': [1,2,3],
    'B': [4,5,6],
    'C': [7,8,9],
    'D': [10,11,12]
})

# 把列C移到最前面
df_relocated_first = relocate(df, 'C', position='first')
print(df_relocated_first.columns)
# 输出: Index(['C', 'A', 'B', 'D'], dtype='object')

# 把列A移到最后面
df_relocated_last = relocate(df, 'A', position='last')
print(df_relocated_last.columns)
# 输出: Index(['B', 'C', 'D', 'A'], dtype='object')

# 把列B移到列D前面
df_relocated_before = relocate(df, 'B', position='before', ref_col='D')
print(df_relocated_before.columns)
# 输出: Index(['A', 'C', 'B', 'D'], dtype='object')

# 把列D移到列A后面
df_relocated_after = relocate(df, 'D', position='after', ref_col='A')
print(df_relocated_after.columns)
# 输出: Index(['A', 'D', 'B', 'C'], dtype='object')

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