如何将一个 numpy 列表数组传递给函数（花式索引）

Question

我想用 numpy 以矢量化的方式使用这个函数：

def example(testing, index):
    return np.sum(testing[index], axis = 1)

假设我们创建了我们的测试数组和索引数组：

test = np.arange(0, 20)
indices = np.array([[0, 2], [1, 3], [0, 3], [1, 2], [3, 4]])

它完成了我的预期，为数组中的每个列表将列表中带有索引的元素相加：

Input: example(test, indices)
Output: [2 4 3 3 7]

但是，如果我尝试使用可变长度的索引列表，

indices = np.array([[0, 2, 3], [1, 3], [0, 3], [1, 2], [3, 4]])
Input: example(test, indices)
Output: IndexError: arrays used as indices must be of integer (or boolean) type

如果不遍历 indices 数组，我无法让 numpy 执行此功能。我知道 numpy 在第一种情况下创建了一个二维数组，但不是在第二种情况下，但我不确定为什么它无法对一维数组执行矢量化元素操作。由于这些数组在现实生活中实际上非常大并且用于并行化函数，因此我想像第一个示例一样以 numpy 风格的矢量化方式执行此操作。

Answer 1

indices 的第一个版本确实是一个二维整数数组 - indices.shape 返回 (5, 2)，print(indices) 给出：

array([[0, 2],
       [1, 3],
       [0, 3],
       [1, 2],
       [3, 4]])

然而，第二个是列表对象的一维数组：

array([list([0, 2, 3]), list([1, 3]), list([0, 3]), list([1, 2]),
       list([3, 4])], dtype=object)

这里根本不需要向量化，因为这些列表不是 numpy 向量。我相信你可以做的最好的矢量化操作是将indices创建为数组列表而不是列表数组：

indices = [np.array(x) for x in [[0, 2, 3], [1, 3], [0, 3], [1, 2], [3, 4]]]

然后重新定义函数来处理这个问题：

def example(testing, index):
    return [ np.sum(testing[i]) for i in index ]

print(example(test, indices)) 
# prints [5, 4, 3, 3, 7]