将带有坐标的一维数组转换为numpy中的二维数组

Convert 1D array with coordinates into 2D array in numpy

我有一个形状为 (N,) 的值数组 arr 和一个形状为 (N,2) 的坐标数组 coords。我想在 (M,M) 数组 grid 中表示它,使得 grid 在不在 coords 中的坐标处取值 0,对于包含的坐标,它应该将具有该坐标的所有值的总和存储在 arr 中。所以如果 M=3、arr = np.arange(4)+1coords = np.array([[0,0,1,2],[0,0,2,2]]) 那么 grid 应该是:

array([[3., 0., 0.],
       [0., 0., 3.],
       [0., 0., 4.]])

之所以如此重要,是因为我需要能够多次重复此步骤,并且 arr 中的值每次都会更改,坐标也可能会更改。理想情况下,我正在寻找矢量化解决方案。我怀疑我可能能够以某种方式使用 np.where,但如何使用并不是很明显。

为解决方案计时

我对此时出现的解决方案进行了计时,累加器方法似乎比稀疏矩阵方法稍快,第二种累加方法最慢,原因在评论中解释:

%timeit for x in range(100): accumulate_arr(np.random.randint(100,size=(2,10000)),np.random.normal(0,1,10000))
%timeit for x in range(100): accumulate_arr_v2(np.random.randint(100,size=(2,10000)),np.random.normal(0,1,10000))
%timeit for x in range(100): sparse.coo_matrix((np.random.normal(0,1,10000),np.random.randint(100,size=(2,10000))),(100,100)).A
47.3 ms ± 1.79 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
103 ms ± 255 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
48.2 ms ± 36 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

一种方法是创建一个 sparse.coo_matrix 并将其转换为密集的:

from scipy import sparse
sparse.coo_matrix((arr,coords),(M,M)).A
# array([[3, 0, 0],
#        [0, 0, 3],
#        [0, 0, 4]])

np.bincount-

def accumulate_arr(coords, arr):
    # Get output array shape
    m,n = coords.max(1)+1

    # Get linear indices to be used as IDs with bincount
    lidx = np.ravel_multi_index(coords, (m,n))
    # Or lidx = coords[0]*(coords[1].max()+1) + coords[1]

    # Accumulate arr with IDs from lidx
    return np.bincount(lidx,arr,minlength=m*n).reshape(m,n)

样本运行-

In [58]: arr
Out[58]: array([1, 2, 3, 4])

In [59]: coords
Out[59]: 
array([[0, 0, 1, 2],
       [0, 0, 2, 2]])

In [60]: accumulate_arr(coords, arr)
Out[60]: 
array([[3., 0., 0.],
       [0., 0., 3.],
       [0., 0., 4.]])

另一个与 np.add.at 相似,可能更容易理解 -

def accumulate_arr_v2(coords, arr):
    m,n = coords.max(1)+1
    out = np.zeros((m,n), dtype=arr.dtype)
    np.add.at(out, tuple(coords), arr)
    return out