如何基于DataFrame元素生成(1, x, x²)形式的三维numpy数组？

先明确原始DataFrame的结构：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'c0': [10, 8, 4],
    'c1': [2, 2, 1]
})

需求很清晰：把每个元素转成[1, x, x²]的数组，最终输出维度为(2, 3, 3)的numpy数组——2对应原DataFrame的两列，3对应三行，最后一个3是每个元素的三个计算值。

为什么你的代码不行？ #

你之前用的df.loc[:,i].apply(lambda x: np.power(x,np.arange(3)))，会返回一个Series，每个元素是[1,x,x²]一维数组，但直接把这个Series转成numpy数组后，得到的是(3,3)的结构（对应一列的3行）。如果遍历所有列拼接，没调整顺序的话维度会是(3,2,3)，和你要的(2,3,3)不符，而且没有自动堆叠成符合要求的三维结构。

两种可行的解决方法 #

方法一：用numpy广播直接计算（高效首选） #

利用numpy的广播特性，一步到位生成目标数组，效率远高于Python循环，适合大数据量场景：

# 把DataFrame转成(3,2)的numpy数组（行×列）
base_arr = df.to_numpy()
# 给数组增加一个维度，配合幂次数组[0,1,2]做广播计算，得到(3,2,3)的中间结果
temp = base_arr[..., np.newaxis] ** np.arange(3)
# 交换前两个维度，把列数放到第一个位置，得到(2,3,3)的最终结果
final_arr = np.transpose(temp, axes=(1, 0, 2))

# 验证维度
print(final_arr.shape)  # 输出 (2, 3, 3)
# 查看第一列（c0）的结果
print(final_arr[0])
# 输出：
# [[  1  10 100]
#  [  1   8  64]
#  [  1   4  16]]

方法二：按列处理后堆叠（直观易懂） #

如果习惯用pandas的apply方法，可以遍历每列处理后再堆叠成目标数组：

col_results = []
for col_name in df.columns:
    # 对当前列的每个元素生成[1,x,x²]，转成(3,3)的数组
    col_data = df[col_name].apply(lambda x: np.power(x, np.arange(3))).to_numpy()
    col_results.append(col_data)
# 把所有列的结果堆叠成(2,3,3)的数组
final_arr = np.stack(col_results)

print(final_arr.shape)  # 输出 (2, 3, 3)

核心要点 #

• 方法一的广播机制避免了Python循环，处理大数据量时速度优势明显
• 最终维度的调整关键是：原始数组是(行,列)，计算幂次后得到(行,列,3)，通过transpose(1,0,2)把列维度移到最前面，就得到了(列,行,3)的目标结构

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