如何将Pandas透视结果转换为三维DataFrame？

透视后得到一个501行×505列的DataFrame，希望将其转换为501行×101列×5深度的三维结构，把x-axis和y-axis对应的不同值（value、value_2…value_5）存储在深度维度中。要求可以通过x-axis和y-axis获取所有对应值（例如x-axis=15、y-axis=20时返回[value, value_2,…,value_5]），理想状态下能通过x-axis、y-axis和值索引（例如15、20、2对应value_3）进行取值。

import pandas as pd 
import numpy as np 
from itertools import product

x_axis = np.arange(0, 501)
y_axis = np.arange(0, 101)

combinations = list(product(x_axis, y_axis))

np.random.seed(42) 
values = np.random.randint(0, 11, len(combinations)) 
values_2 = np.random.randint(0, 11, len(combinations)) 
values_3 = np.random.randint(0, 11, len(combinations)) 
values_4 = np.random.randint(0, 11, len(combinations)) 
values_5 = np.random.randint(0, 11, len(combinations))

df = pd.DataFrame(combinations, columns=['x-axis', 'y-axis'])
df['value'] = values 
df['value_2'] = values_2 
df['value_3'] = values_3 
df['value_4'] = values_4 
df['value_5'] = values_5

pivot_df = df.pivot(index='x-axis', columns='y-axis')

方法一：NumPy三维数组（纯数值索引） #

利用pivot后DataFrame的多层列结构，通过NumPy重塑维度并调整顺序，得到符合要求的三维数组：

# 提取pivot后的数值矩阵，形状为(501, 505)
pivot_values = pivot_df.values
# 先重塑为(501, 5, 101)，对应[x轴, 值类型, y轴]
reshaped_matrix = pivot_values.reshape(501, 5, 101)
# 调整维度顺序为(501, 101, 5)，对应[x轴, y轴, 值深度]
three_d_array = reshaped_matrix.transpose(0, 2, 1)

# 取值示例1：x=15、y=20时获取所有对应值
print(three_d_array[15, 20])  # 输出格式：[value, value_2, value_3, value_4, value_5]
# 取值示例2：x=15、y=20、索引2对应value_3
print(three_d_array[15, 20, 2])

方法二：xarray带标签索引（更直观） #

如果需要保留x-axis、y-axis、值名称的标签进行索引，推荐使用xarray库，它原生支持多维带标签的数据操作：

import xarray as xr

# 将pivot后的DataFrame转换为xarray数据集
ds = xr.Dataset.from_dataframe(pivot_df)
# 转换为三维DataArray，调整维度顺序为[x-axis, y-axis, 值类型]
three_d_da = ds.to_array(dim='value_type').transpose('x-axis', 'y-axis', 'value_type')

# 取值示例1：通过标签获取x=15、y=20的所有值
print(three_d_da.sel({'x-axis':15, 'y-axis':20}).values)
# 取值示例2：通过标签直接获取x=15、y=20的value_3
print(three_d_da.sel({'x-axis':15, 'y-axis':20, 'value_type':'value_3'}).values)

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