将数组值快速求和到python中的矩阵中

Question

我有两个带索引（行和列）的数组和一个带值的数组。我想做的是使用具有行和列索引的数组将数组中的值添加到相应矩阵中的值中。我还需要记录我在某个地方添加了多少次价值。

使用 for 循环没有问题。然而，在我的情况下，这将需要很多小时才能完成。 那么有没有更快的方法来完成下面的代码呢？

number = np.zeros([2,2])
amount = np.zeros([2,2])
row = np.array([1, 1, 1, 0, 1, 1, 0, 0, 1, 0])
col = np.array([1, 0, 1, 0, 1, 1, 1, 1, 0, 1])
values = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
for i in range(len(row)):
    r = row[i]
    c = col[i]
    v = values[i]
    amount[r, c] = amount[r, c] + v
    number[r, c] = number[r, c] + 1

此示例中的结果将如下所示

amount = array([[  4.,  25.],
                [ 11.,  15.]])

和

number = array([[ 1.,  3.],
                [ 2.,  4.]])

Answer 1

如您所见，numpy 不适合这类问题，至少在我看来不是。这是我可能用 cython 或 C 编写的东西，并与我的其余 python 代码粘合在一起。

一个简单的迭代解决方案将花费values 大小的线性时间。我能想到的使用 numpy 的最佳解决方案需要排序，因此需要 O(n log n)。对于少量输入，它仍然会比纯 Python 中的迭代解决方案更快，但渐近实际上应该更糟。

无论如何，代码如下：

def numpy_assign(values, row, col):
    # arrange the values so that entries are contiguous
    ix = np.ravel_multi_index((row, col), (2,2))
    sort_ix = np.argsort(ix)

    ix = ix[sort_ix]
    values = values[sort_ix]

    # now we need to construct the funny index argument 
    # to the `reduceat` ufunc method
    reduce_ix = np.where(np.diff(ix))[0] + 1
    reduce_ix = np.insert(reduce_ix, 0, 0)

    # reduce the amounts
    amount = np.add.reduceat(values, reduce_ix).reshape(2,2)

    # and count the number of times we've done so
    _, counts = np.unique(ix, return_counts=True)
    counts = counts.reshape(2,2)

    return amount, counts

现在，如果您的 python 发行版中有 numba（我只是使用 Anaconda 发行版的学术许可证附带的标准安装），那么我们实际上可以做得更好而不需要太多额外的工作：

import numba

@numba.autojit
def inner_loop(row, col, values, number, amount):
    for i in range(len(row)):
        r = row[i]
        c = col[i]
        v = values[i]
        amount[r, c] += v
        number[r, c] += 1

def numba_assign(row, col, values):
    number = np.zeros([2,2], np.float)
    amount = np.zeros([2,2], np.float)
    inner_loop(row, col, values, number, amount)
    return amount, number

这能快多少？在 1e6 个条目的随机输入上，我们得到时间：

>>> %timeit loop_assign(row, col, values)
1 loops, best of 3: 1.35 s per loop

>>> %timeit numpy_assign(row, col, values)
10 loops, best of 3: 81 ms per loop

>>> %timeit numba_assign(row, col, values)
100 loops, best of 3: 5.33 ms per loop

loop_assign 是您的原始代码。请注意，numba_assign 不仅运行速度比 numpy 解决方案快得多，而且它不需要排序，因此它具有线性渐近复杂度，这意味着它可以更好地扩展。