如何基于索引规则实现Pandas DataFrame的广播展开？

问题场景 #

现有如下DataFrame，其索引为包含数组的Index对象：

import pandas as pd 

df = pd.DataFrame([[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9]],
                  index=pd.Index([[1], [2, 3], [4]]),
                  columns=['a', 'b', 'c'])
print(df)

需要将索引中的数组展开，对应行自动广播复制，最终得到如下结果：

expected = pd.DataFrame([[1, 2, 3],
                         [4, 5, 6],
                         [4, 5, 6],
                         [7, 8, 9]],
                        index=[1, 2, 3, 4],
                        columns=['a', 'b', 'c'])
print(expected)    

解决方案 #

方法1：兼容低版本Pandas的通用写法

如果你的Pandas版本低于1.3.0，可通过「索引转列→展开→重置索引」的步骤实现：

# 将原索引转为普通列，命名为'idx'
df_reset = df.reset_index(names='idx')
# 展开数组列，自动复制对应行
df_exploded = df_reset.explode('idx')
# 将展开后的列重新设为索引
result = df_exploded.set_index('idx')

print(result)

方法2：高版本Pandas简洁写法

Pandas 1.3.0及以上版本支持直接对索引执行explode操作，一行代码即可完成：

result = df.explode(index=True)
print(result)

两种方法最终输出都与预期结果一致。

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