花哨索引的Numpy循环广播

时间:2017-07-21 23:59:25

标签: python numpy broadcast

Anumpy数组,其形状为(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_tostride_tricks是否比嵌套的python循环更快。 (Repeat NumPy array without replicating data?

1 个答案:

答案 0 :(得分:2)

您可以将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.]])