Python 加速使用切片将 pandas DataFrame 重组为 numpy 数组的循环会抛出异常 我错过了什么

我有一个Pandas 数据框，是这样的:

raw_data = DataFrame({
    'date_idx': [0, 1, 2, 0, 1, 2],
    'element_idx': [0, 0, 0, 1, 1, 1],
    'a': [10, 20, 30, 40, 50, 60],
    'b': [11, 21, 31, 41, 51, 61],
    'c': [12, 22, 32, 42, 52, 62],
})

我将date_idx和element_idx之外的列称为"输入".我想用date_idx-&gt；input_idx-&gt；element_idx将它重新组织成一个3D Numpy数组，结果如下:

[[[10. 40.]
  [11. 41.]
  [12. 42.]]

 [[20. 50.]
  [21. 51.]
  [22. 52.]]

 [[30. 60.]
  [31. 61.]
  [32. 62.]]]

我使用了两个for循环，它运行得很好:

date_idx = [0, 1, 2, 0, 1, 2]
element_idx = [0, 0, 0, 1, 1, 1]
raw_data = DataFrame({
    'date_idx': date_idx,
    'element_idx': element_idx,
    'a': [10.0, 20.0, 30.0, 40.0, 50.0, 60.0],
    'b': [11.0, 21.0, 31.0, 41.0, 51.0, 61.0],
    'c': [12.0, 22.0, 32.0, 42.0, 52.0, 62.0],
})

inputs = ['a', 'b', 'c']

unique_dates = set(date_idx)
unique_elements = set(element_idx)
data = np.zeros(shape=(len(unique_dates), len(inputs), len(unique_elements)), dtype=np.float64)

for i in range(len(raw_data)):
    row = raw_data.iloc[i]
    date_idx = int(row['date_idx'])
    element_idx = int(row['element_idx'])

    for input_idx in range(len(inputs)):
        data[date_idx][input_idx][element_idx] = float(row[inputs[input_idx]])

print(data)

然而，这是非常缓慢的.对于date_idx数组，我有数百万个条目，inputs和element_idx都有几十个条目.在我的机器上花了7个小时才用我的真实数据集完成了这项工作.

我有一种感觉，这可以通过切片来完成，没有循环，但我的try 总是失败-我错过了一些东西.

例如，我try 用以下命令消除内部循环:

for i in range(len(raw_data)):
    row = raw_data.iloc[i]
    date_idx = int(row['date_idx'])
    element_idx = int(row['element_idx'])

    data[date_idx][:][element_idx] = list(dict(row[inputs]).values())

它失败了，原因是:

Traceback (most recent call last):
  File "/home/stark/Work/mmr6/test2.py", line 84, in <module>
    data[date_idx][:][element_idx] = list(dict(row[inputs]).values())
    ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^
ValueError: could not broadcast input array from shape (3,) into shape (2,)

我问题是，是否可以使用切片和/或快速技术在普通的NumPy数组上以这种方式重新组织这个DataFrame，或者我真的需要这里的循环吗？