如何在Python中编写支持任意数量输入的DataFrame公共索引函数？

问题描述 #

我希望编写一个保留输入DataFrame公共索引的函数。针对2个DataFrame，我已实现如下代码：

def get_df_common_index(df1, df2):
    common_index = df1.index.intersection(df2.index)
    return df1.loc[common_index], df2.loc[common_index]

现需将其改造为支持任意数量输入的通用函数，使其可按如下方式调用并返回对应DataFrame：

df1, df2, df3, df4 = get_df_common_index(df1, df2, df3, df4)

请问如何在Python中实现这类支持未知数量输入的函数？

解决方案 #

在Python中，可以通过*args语法接收任意数量的位置参数，以此实现支持未知数量DataFrame输入的函数，具体实现如下：

def get_df_common_index(*dfs):
    # 处理无输入的边界情况
    if not dfs:
        return ()
    # 初始化公共索引为第一个DataFrame的索引
    common_index = dfs[0].index
    # 遍历所有剩余DataFrame，逐步计算索引交集
    for df in dfs[1:]:
        common_index = common_index.intersection(df.index)
        # 若中途公共索引为空，提前终止循环优化性能
        if common_index.empty:
            break
    # 对每个输入DataFrame筛选公共索引，返回结果元组
    return tuple(df.loc[common_index] for df in dfs)

关键说明 #

• *dfs参数会将所有传入的位置参数打包成一个元组，让函数可以接收任意数量的DataFrame输入。
• 以第一个DataFrame的索引为初始值，依次与后续每个DataFrame的索引求交集，最终得到所有输入的公共索引。
• 最后将每个DataFrame按公共索引筛选后的结果打包成元组返回，支持直接解包赋值给多个变量。

测试示例 #

import pandas as pd

# 构造测试用DataFrame
df1 = pd.DataFrame({'col1': [1,2,3]}, index=['a','b','c'])
df2 = pd.DataFrame({'col2': [4,5]}, index=['b','c'])
df3 = pd.DataFrame({'col3': [6]}, index=['c'])
df4 = pd.DataFrame({'col4': [7,8,9]}, index=['a','c','d'])

# 调用函数并解包结果
df1, df2, df3, df4 = get_df_common_index(df1, df2, df3, df4)

print(df1.index)  # 输出：Index(['c'], dtype='object')
print(df2.index)  # 输出：Index(['c'], dtype='object')

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