在索引两个数组时进行numpy矢量化

Question

我正在尝试使用numpy将以下函数向量化，并且完全丢失了。

A = ndarray: Z x 3

B = ndarray: Z x 3

C = integer

D = ndarray: C x 3

伪代码：

entries = []
means = []
For i in range(C):
    for p in range(len(B)):
        if B[p] == D[i]:
            entries.append(A[p])
    means.append(columnwise_means(entries))
return means

一个例子是：

A = [[1,2,3],[1,2,3],[4,5,6],[4,5,6]]
B = [[9,8,7],[7,6,5],[1,2,3],[3,4,5]]
C = 2
D = [[1,2,3],[4,5,6]]

返回值：

[average([9,8,7],[7,6,5]), average(([1,2,3],[3,4,5])] = [[8,7,6],[2,3,4]]

我试过使用np.where，np.argwhere，np.mean等，但是似乎无法获得预期的效果。 任何帮助将不胜感激。

谢谢！

Answer 1

我看到两个提示：

首先，按行比较数组。 一种方法是简化一维索引系统：

def indexer(M,base=256):
    return (M*base**arange(3)).sum(axis=1)

base是一个整数> A.max() 。 然后可以这样选择：

indices=np.equal.outer(indexer(D),indexer(A))

为：

array([[ True,  True, False, False],
       [False, False,  True,  True]], dtype=bool)

其次，每个组的长度可以不同，因此最后一步很难进行向量化。 这是完成工作的一种方法。

B=array(B)
means=[B[i].mean(axis=0) for  i in indices]

Answer 2

按照问题的预期输出，我假设在实际代码中，您将：

• IF条件语句为： if A[p] == D[i] ，并且

• 条目将从B： entries.append(B[p])追加。

因此，这是NumPy broadcasting和dot-product的一种矢量化方法-

mask = (D[:,None,:] == A).all(-1)
out = mask.dot(B)/(mask.sum(1)[:,None])

如果输入数组是整数数组，则可以将内存视为n维数组的索引，从而节省内存并提高性能，从而无需像这样进行3D即可创建2D mask -

dims = np.maximum(A.max(0),D.max(0))+1
mask = np.ravel_multi_index(D.T,dims)[:,None] == np.ravel_multi_index(A.T,dims)

样品运行-

In [107]: A
Out[107]: 
array([[1, 2, 3],
       [1, 2, 3],
       [4, 5, 6],
       [4, 5, 6]])

In [108]: B
Out[108]: 
array([[9, 8, 7],
       [7, 6, 5],
       [1, 2, 3],
       [3, 4, 5]])

In [109]: mask = (D[:,None,:] == A).all(-1)
     ...: out = mask.dot(B)/(mask.sum(1)[:,None])
     ...: 

In [110]: out
Out[110]: 
array([[8, 7, 6],
       [2, 3, 4]])