给定 numpy 中的索引列表，添加元素的最有效方法

Question

假设我们有一个形状为 (N, ) 的 numpy 数组 A 和一个形状为 (M, 3) 的矩阵 D，其中包含数据和另一个形状为 (M, 3) 的矩阵 I，它具有 D 中每个数据元素的相应索引。如何我们可以构造 A 给定 D 和 I 以便添加重复的元素索引吗？

例子：

############# A[I] := D ###################################  
A = [0.5, 0.6]                         # Final Reduced Data Vector
D = [[0.1, 0.1 0.2], [0.2, 0.4, 0.1]]  # Data
I = [[0, 1, 0], [0, 1, 1]]             # Indices

例如：

A[0] = D[0][0] + D[0][2] + D[1][0]     # 0.5 = 0.1 + 0.2 + 0.2

由于在索引矩阵中我们有：

I[0][0] = I[0][2] = I[1][0] = 0

目标是避免循环遍历所有元素以对大 N、M (10^6-10^9) 有效。

Answer 1

我怀疑你能比np.bincount - 请注意官方文档如何提供这个确切的用例

# Your example
A = [0.5, 0.6]
D = [[0.1, 0.1, 0.2], [0.2, 0.4, 0.1]]
I = [[0, 1, 0], [0, 1, 1]]

# Solution
import numpy as np    
D, I = np.array(D).flatten(), np.array(I).flatten()
print(np.bincount(I, D)) #[0.5 0.6]

Answer 2

I和D的形状无关紧要：您可以清楚地解开 arrays 而不改变结果：

index = np.ravel(I)
data = np.ravel(D)

现在您可以根据I对 arrays 进行排序：

sorter = np.argsort(index)
index = index[sorter]
data = data[sorter]

这很有帮助，因为现在index看起来像这样：

0, 0, 0, 1, 1, 1

data是这样的：

0.1, 0.2, 0.2, 0.1, 0.4, 0.1

将连续数字的运行加在一起应该比处理随机位置更容易。 让我们首先找到运行开始的索引：

runs = np.r_[0, np.flatnonzero(np.diff(index)) + 1]

现在您可以使用像 np.add 这样的np.add具有称为reduceat的部分reduce操作这一事实。 这允许您对数组的区域求和：

a = np.add.reduceat(data, runs)

如果保证I至少包含一次 [0, A.size ) 中的所有索引，那么您就完成了：只需分配给A而不是a 。 如果没有，您可以使用index中每次运行的开始是目标索引的事实进行映射：

A = np.zeros(n)
A[index[runs]] = a

算法复杂度分析：

• 如果数据在数组中，则ravel在时间和空间上为 O(1)。 如果是一个列表，这在时间和空间上都是 O(MN)
• argsort在时间上是 O(MN log MN)，在空间上是O(MN)
• sorter的索引在时间和空间上是 O(MN)
• 计算runs时间为 O(MN)，空间为 O(MN + M) = O(MN)
• reduceat是单程：时间 O(MN)，空间 O(M)
• 重新分配A在时间和空间上是 O(M)

总计：O(MN log MN)时间，O(MN)空间

TL;博士

def make_A(D, I, M):
    index = np.ravel(I)
    data = np.ravel(D)
    sorter = np.argsort(index)
    index = index[sorter]

    if index[0] < 0 or index[-1] >= M:
        raise ValueError('Bad indices')

    data = data[sorter]
    runs = np.r_[0, np.flatnonzero(np.diff(index)) + 1]
    a = np.add.reduceat(data, runs)
    if a.size == M:
        return a
    A = np.zeros(M)
    A[index[runs]] = a
    return A

Answer 3

如果您事先知道 A 的大小，就像您所做的那样，您可以简单地使用add.at ：

import numpy as np

D = [[0.1, 0.1, 0.2], [0.2, 0.4, 0.1]]
I = [[0, 1, 0], [0, 1, 1]]

arr_D = np.array(D)
arr_I = np.array(I)

A = np.zeros(2)

np.add.at(A, arr_I, arr_D)

print(A)

Output

[0.5 0.6]

如果你不知道 A 的大小，你可以使用max来计算它：

A = np.zeros(arr_I.max() + 1)
np.add.at(A, arr_I, arr_D)
print(A)

Output

[0.5 0.6]

该算法的时间复杂度为O(N) ，空间复杂度为O(N) 。

这：

arr_I.max() + 1

是 bincount 在幕后所做的，来自文档：

对输入数组进行分箱的结果。 out 的长度等于 np.amax(x)+1。

话虽如此， bincount 至少要快一个数量级：

I = np.random.choice(1000, size=(1000, 3), replace=True)
D = np.random.random((1000, 3))
%timeit make_A_with_at(I, D, 1000)
213 µs ± 25 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit make_A_with_bincount(I, D)
11 µs ± 15.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)