在索引到数组时消除for循环

Question

我有两个数组：vals具有形状(N，m(，其中N约为100万，m为3。这些值是浮点值我有另一个形状为(N,4)的阵列indices。indices中的所有值都是vals中的行索引。(此外，与这里的示例不同，indices的每一行都包含唯一的值。(。

import numpy as np
from random import randrange
# set up the arrays for this test example (no need to improve this)
N = 9
vals = np.array(list(range(3*N))).reshape((N,3))
indices = np.array([randrange(N) for n in range(4*N)]).reshape((N,4))

创建阵列aug时，我想替换以下for循环

# form an augmented matrix by indexing into vals using rows from indices
aug = np.stack([vals[indices[x]] for x in range(N)])
# compute a mean along axis=1 of aug
aug.mean(axis=1)

该问题的更广泛上下文是vals包含三维分布粒子的数值数据。使用对粒子的空间位置的最近邻居搜索(使用scipy.spatial.cKDTree(来生成CCD_ 9。我想对最近邻居的数字数据求平均值。由于我有大约100万个粒子，for循环相当慢。

Answer 1

您实际上可以用替换整个aug = ...行

aug = vals[indices]

这将产生相同的结果：

np.array_equal(
np.stack([vals[indices[x]] for x in range(N)]),
vals[indices]
)
# True