花式索引的 Numpy 循环广播

Question

A 是一个 numpy 数组，形状为 (6, 8)

我要：

x_id = np.array([0, 3])
y_id = np.array([1, 3, 4, 7])

A[ [x_id, y_id] += 1  # this doesn't actually work.

像::2这样的技巧不会奏效，因为索引不会定期增加。

我不想使用额外的内存来重复 [0, 3] 并创建一个新数组 [0, 3, 0, 3]，因为那样很慢。

两个维度的索引长度不相等。

相当于：

A[0, 1] += 1
A[3, 3] += 1
A[0, 4] += 1
A[3, 7] += 1

可以numpy做这样的事情吗？

更新：

不确定 broadcast_to 或 stride_tricks 是否比嵌套 python 循环更快。 (Repeat NumPy array without replicating data?)

Answer 1

你可以把y_id转成一个二维数组，二维数组和x_id一样，然后两个索引会因为维度不同而自动广播：

x_id = np.array([0, 3])
y_id = np.array([1, 3, 4, 7])

A = np.zeros((6,8))
A[x_id, y_id.reshape(-1, x_id.size)] += 1 

A
array([[ 0.,  1.,  0.,  0.,  1.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  1.,  0.,  0.,  0.,  1.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.]])

花式索引的 Numpy 循环广播

Numpy Cyclic Broadcast of Fancy Indexing

python

numpy

broadcast