以向量化方式访问 numpy ndarray 的列

Question

我有一个索引列表列表，如下所示：

outer_list = [[0], [1, 2], [1, 2], [3, 5], [4], [3, 5]]

每个内部列表包含其自身在外部列表中的索引号。

我还有一个带有值的 numpy 数组

val_mat

。以两行为例：

val_mat = np.arange(2*6).reshape(2, 6)

我想获取每个索引组的行组平均值并将其存储在平均值矩阵中。在此示例中，我希望（在每一行中）索引 0 处的元素具有其值，索引 1 和 2 处的元素具有索引 1 和 2 处的值的行平均值，以及索引 3 处的元素和 5 得到元素 3 和 5 的平均值，等等

目前我正在使用：

mean_vals = np.vstack([np.nanmean(val_mat[:, inner_list], axis=1) for inner_list in outer_list]).T

但是，外部列表非常大（500,000 个内部列表），在 10k x 500k 矩阵中生成所有平均值大约需要 2 分钟。

有没有一种方法可以以矢量化方式执行此操作，而无需迭代外部列表？

Answer 1

所以你的结果是

In [93]: mean_vals
Out[93]: 
array([[ 0. ,  1.5,  1.5,  4. ,  4. ,  4. ],
       [ 6. ,  7.5,  7.5, 10. , 10. , 10. ]])

如果我更改

outer_list

，那么所有内部列表都具有相同的长度：

In [94]: outer_list = [[0,0], [1, 2], [1, 2], [3, 5], [4,4], [3, 5]]

我在没有列表理解的情况下得到相同的平均值：

In [95]: np.mean(val_mat[:, outer_list], axis=2)
Out[95]: 
array([[ 0. ,  1.5,  1.5,  4. ,  4. ,  4. ],
       [ 6. ,  7.5,  7.5, 10. , 10. , 10. ]])

In [96]: val_mat[:, outer_list]
Out[96]: 
array([[[ 0,  0],
        [ 1,  2],
        [ 1,  2],
        [ 3,  5],
        [ 4,  4],
        [ 3,  5]],

       [[ 6,  6],
        [ 7,  8],
        [ 7,  8],
        [ 9, 11],
        [10, 10],
        [ 9, 11]]])