如何将分组Pandas DataFrame转换为形状(4,2,3)的3D Numpy数组？

尝试将分组后的Pandas DataFrame转换为3D Numpy数组，要求每个DataFrame单元格映射为1D Numpy数组，期望得到形状为(4, 2, 3)的3D数组，但实际得到形状为(4, 2)的2D数组，其中每个元素是形状为(3,)的1D数组。

简化示例代码 #

import pandas as pd
import numpy as np

example_df = pd.DataFrame(
    [["A", True, True], 
     ["A", True, False], 
     ["B", False, True], 
     ["B", False, False]],
    columns=["grp", "col1", "col2"],
)


def mapper(val):
    return np.ones(3) if val is True else np.zeros(3)


result = (
    example_df.groupby(["grp"])
    .apply(lambda df: df.map(mapper), include_groups=False)
    .to_numpy()
)

print(type(result))
print(type(result[0]))
print(type(result[0][0]))
print(type(result[0][0][0]))
print(result.shape)
print(result[0, 0].shape)

实际输出 #

<class 'numpy.ndarray'>
<class 'numpy.ndarray'>
<class 'numpy.ndarray'>
<class 'numpy.float64'>
(4, 2)
(3,)

出现问题的核心原因是：df.map(mapper)返回的DataFrame中每个元素都是独立的Numpy数组，调用to_numpy()时Pandas会将其转换为object类型的2D数组，而非自动堆叠成3D数组。以下是几种可行的解决方法：

方法1：用np.stack二次处理结果 #

先按原有逻辑生成2D数组，再通过np.stack将元素沿第三维度堆叠：

# 沿用原代码生成temp_result
temp_result = (
    example_df.groupby(["grp"])
    .apply(lambda df: df.map(mapper), include_groups=False)
    .to_numpy()
)

# 堆叠为3D数组
result = np.stack(temp_result.flatten()).reshape(temp_result.shape + (3,))
print(result.shape)  # 输出 (4, 2, 3)

方法2：在分组处理中直接生成数组 #

跳过中间的DataFrame转换，在分组后的apply里直接处理并堆叠数组：

result = np.vstack([
    np.column_stack([mapper(val) for val in df[col]] for col in ["col1", "col2"])
    for _, df in example_df.groupby("grp")
])
print(result.shape)  # 输出 (4, 2, 3)

方法3：向量化处理（效率最优） #

如果映射逻辑可以向量化，直接对布尔数组进行扩展，避免循环：

# 将布尔列转换为数组，扩展第三维度并重复3次
bool_array = example_df[["col1", "col2"]].to_numpy()
result = np.repeat(bool_array[:, :, np.newaxis], 3, axis=2).astype(np.float64)
print(result.shape)  # 输出 (4, 2, 3)

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