Pandas DataFrame 到多维 NumPy 数组

Question

我有一个数据框，我想使用其中一列作为第三维将其转换为多维数组。
举个例子：

df = pd.DataFrame({
'id': [1, 2, 2, 3, 3, 3],
'date': np.random.randint(1, 6, 6),
'value1': [11, 12, 13, 14, 15, 16],
'value2': [21, 22, 23, 24, 25, 26]
 })

我想将其转换为具有如下维度（id、日期、值）的 3D 数组：

问题是 'id' 的出现次数不同，所以我不能使用np.reshape()。

对于这个简化的例子，我可以使用：

ra = np.full((3, 3, 3), np.nan)

for i, value in enumerate(df['id'].unique()):
    rows = df.loc[df['id'] == value].shape[0]
    ra[i, :rows, :] = df.loc[df['id'] == value, 'date':'value2']

要产生所需的结果：

但原始 DataFrame 包含数百万行。

是否有一种矢量化的方式来实现相同的结果？

Answer 1

方法#1

这是 @Yannis 在 cmets 中建议的将 id col 与 df.sort_values('id', inplace=True) 排序后的一种矢量化方法 -

count_id = df.id.value_counts().sort_index().values
mask = count_id[:,None] > np.arange(count_id.max())
vals = df.loc[:, 'date':'value2'].values
out_shp = mask.shape + (vals.shape[1],)
out = np.full(out_shp, np.nan)
out[mask] = vals

方法#2

另一个 factorize 不需要任何预排序 -

x = df.id.factorize()[0]   
y = df.groupby(x).cumcount().values
vals = df.loc[:, 'date':'value2'].values
out_shp = (x.max()+1, y.max()+1, vals.shape[1])
out = np.full(out_shp, np.nan)
out[x,y] = vals