如何将二级pd.MultiIndex转换为指定形状的3D NumPy数组

问题场景 #

你有一个带instances和timepoints二级索引的DataFrame：

var_0  var_1
instances timepoints              
0         1               1      4
          2               2      5
          3               3      6
          4               5      8
1         1               1      4
          2               2     55
          3               3      6
          4               3      6
2         1               1     42
          2               2      5
          3               3      6

需要将其转换为形状为(n_instances, n_columns, n_timepoints)的三维NumPy数组，但现有代码生成了object类型的数组，无法得到规整的三维结构。

现有代码问题 #

你的代码使用df_train.index.levels[0]获取实例值，但levels包含索引的所有可能值（可能存在无数据的实例），且不同实例的时间点数量不一致时，列表推导生成的数组会是object dtype，无法形成规整的三维数组。

解决方案 #

根据数据中时间点的完整性，分两种情况处理：

情况1：所有实例的时间点数量一致

直接按instances分组，提取每组数据后转成数组，再调整维度顺序：

# 按instances分组，提取每组的数值
grouped = [group.values for _, group in df_train.groupby(level='instances')]
# 转换为三维数组，并调整维度为(n_instances, n_columns, n_timepoints)
result = np.array(grouped).transpose(0, 2, 1)
print(result.shape)  # 输出：(3, 2, 4)

情况2：实例的时间点数量不一致

先将timepoints转成列（补全缺失时间点为NaN），再重塑形状：

# 将timepoints索引转成列，缺失的时间点自动填充NaN
unstacked_df = df_train.unstack(level='timepoints')
# 重塑为目标形状：(实例数, 列数, 时间点数)
result = unstacked_df.values.reshape(
    unstacked_df.shape[0],       # n_instances
    df_train.shape[1],           # n_columns
    unstacked_df.shape[1] // df_train.shape[1]  # n_timepoints
)
print(result.shape)  # 输出：(3, 2, 4)，缺失的时间点对应值为NaN

优化你的原有代码

如果要保留原有思路，改用index.unique(level='instances')获取实际存在的实例，并处理形状不一致的情况：

# 获取实际存在的实例（避免levels中无数据的实例）
unique_cases = df_train.index.unique(level='instances')
# 提取每个实例的数值数组
data_list = [df_train.loc[case].values for case in unique_cases]

# 检查所有实例的时间点数量是否一致
if all(arr.shape == data_list[0].shape for arr in data_list):
    result = np.array(data_list).transpose(0, 2, 1)
else:
    # 时间点数量不一致时，用unstack方法补全
    unstacked_df = df_train.unstack(level='timepoints')
    result = unstacked_df.values.reshape(unstacked_df.shape[0], df_train.shape[1], -1)

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