花式索引的 Numpy 循环广播

Posted 2023-03-11

【中文标题】花式索引的 Numpy 循环广播

【英文标题】：Numpy Cyclic Broadcast of Fancy Indexing

【发布时间】：2017-12-28 03:14:36

【问题描述】：

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?)

【问题讨论】：

希望存在：***.com/questions/5564098/…

【参考方案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.]])

【讨论】：

听起来很棒...实际的数组是 6 维的。让我试试看。