如何将Pandas DataFrame转换为指定形状的3D NumPy矩阵？

问题描述 #

给定以下DataFrame（可通过pd.DataFrame(data_dict)创建）：

data_dict = {
    'Elevation': {0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 1, 6: 1, 7: 1, 8: 1, 9: 1},
    'Azimuth': {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 0, 6: 1, 7: 2, 8: 3, 9: 4},
    'median': {0: 255, 1: 255, 2: 255, 3: 255, 4: 255, 5: 256, 6: 256, 7: 256, 8: 256, 9: 256},
    'count': {0: 255, 1: 255, 2: 255, 3: 255, 4: 255, 5: 250, 6: 250, 7: 250, 8: 250, 9: 250},
    'to_drop': {0: 1, 1: 1, 2: 1, 3: 1, 4: 1, 5: 0, 6: 0, 7: 0, 8: 0, 9: 0}
}

需要将其转换为形状为(5, 2, 3)的NumPy 3D矩阵：

• 第一维度：Azimuth的唯一值数量（5个）
• 第二维度：Elevation的唯一值数量（2个）
• 第三维度：对应median、count、to_drop三列数据

尝试使用data.groupby(['Elevation','Azimuth']).apply(lambda x: x.values).reset_index().values仅得到(10, 3)的数组，无法满足需求，求可行方法。

解决方案 #

方法1：利用pivot_table重塑后调整形状 #

先通过pivot_table将数据按Azimuth和Elevation重新组织，再转换为数组并调整维度：

import pandas as pd
import numpy as np

# 构建DataFrame
df = pd.DataFrame(data_dict)

# 生成透视表：行=Azimuth，列=Elevation，值为三个目标列
pivot_df = df.pivot_table(
    index='Azimuth',
    columns='Elevation',
    values=['median', 'count', 'to_drop']
)

# 转换为NumPy数组并调整为目标形状(5,2,3)
result = pivot_df.values.reshape(5, 2, 3)

说明：pivot_table生成的是(5,6)的结构（5行Azimuth × 2列Elevation × 3个指标），通过reshape将第二个维度拆分为(2,3)，完美匹配需求。

方法2：分组+展开+形状调整 #

通过groupby分组提取目标数据，再用unstack展开维度，最后调整形状：

# 按Azimuth和Elevation分组，提取三个列的数组
grouped = df.groupby(['Azimuth', 'Elevation'])[['median', 'count', 'to_drop']].apply(np.array)

# 将Elevation维度展开为列，再转换为数组并调整形状
result = grouped.unstack().values.reshape(5, 2, 3)

说明：分组后每个组对应一个长度为3的数组，unstack将Elevation转为列，得到(5,2)的结构，每个元素是长度3的数组，reshape后直接得到目标3D矩阵。

验证结果 #

检查输出形状：

print(result.shape)  # 输出: (5, 2, 3)

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