通过高级索引/向量化更新矩阵

时间:2020-09-02 22:16:01

标签: python arrays numpy vectorization

我有一个有效计数器的矩阵。我想基于列索引列表来增加这些计数器-每个位置索引也对应于一行增量。

使用for循环很简单,而使用列表理解则不太简单。无论哪种情况,都涉及迭代。但是我想知道是否有任何方法可以向量化这个问题?

最小的问题是:

counters = np.zeros((4,4))
counters
array([[0., 0., 0., 0.],
       [0., 0., 0., 0.],
       [0., 0., 0., 0.],
       [0., 0., 0., 0.]])


update_columns = [1,0,2,2]

for row, col in zip(range(len(update_columns)), update_columns):
    counters[row, col] += 1
    
counters
array([[0., 1., 0., 0.],
       [1., 0., 0., 0.],
       [0., 0., 1., 0.],
       [0., 0., 1., 0.]])

1 个答案:

答案 0 :(得分:1)

您正在寻找的被称为高级numpy索引。您可以使用np.arange传递行索引,并使用update_columns传递列索引:

update_columns = np.array(update_columns)
counters[np.arange(update_columns.size), update_columns] += 1

输出:

[[0. 1. 0. 0.]
 [1. 0. 0. 0.]
 [0. 0. 1. 0.]
 [0. 0. 1. 0.]]