问题:向量化 NumPy 数组中的迭代加法

对于 2D 索引的随机数组中的每个元素(具有潜在的重复项),我想“+u003d1”到 2D 零数组中的相应网格。但是,我不知道如何优化计算。使用标准 for 循环,如此处所示,

def interadd():
    U = 100 
    input = np.random.random(size=(5000,2)) * U
    idx = np.floor(input).astype(np.int) 

    grids = np.zeros((U,U))      
    for i in range(len(input)):
        grids[idx[i,0],idx[i,1]] += 1
    return grids

运行时可能非常重要:

>> timeit(interadd, number=5000)
43.69953393936157

有没有办法向量化这个迭代过程?

解答

您可以使用np.add.at来加快速度,它可以正确处理重复索引的情况:

def interadd(U, idx):
    grids = np.zeros((U,U))      
    for i in range(len(idx)):
        grids[idx[i,0],idx[i,1]] += 1
    return grids

def interadd2(U, idx):
    grids = np.zeros((U,U))
    np.add.at(grids, idx.T.tolist(), 1)
    return grids

def interadd3(U, idx):
    # YXD suggestion
    grids = np.zeros((U,U))
    np.add.at(grids, (idx[:,0], idx[:,1]), 1)
    return grids

这使

>>> U = 100
>>> idx = np.floor(np.random.random(size=(5000,2))*U).astype(np.int)
>>> (interadd(U, idx) == interadd2(U, idx)).all()
True
>>> %timeit interadd(U, idx)
100 loops, best of 3: 8.48 ms per loop
>>> %timeit interadd2(U, idx)
100 loops, best of 3: 2.62 ms per loop

以及 YXD 的建议:

>>> (interadd(U, idx) == interadd3(U, idx)).all()
True
>>> %timeit interadd3(U, idx)
1000 loops, best of 3: 1.09 ms per loop