如何利用1D数组索引对Numpy二维数组对应元素执行加1操作？

给定二维Numpy数组A：

import numpy as np

A = np.array([[0, 1, 3, 5, 9], 
              [2, 7, 5, 1, 4]])

和一维数组B：

B = np.array([2, 4])

需求是给A每行中对应B索引位置的元素加1，得到目标数组：

array([[0, 1, 4, 5, 9],
       [2, 7, 5, 1, 5]])

实现方法 #

利用Numpy的整数数组索引特性，直接定位到每行需要修改的位置：

# 生成对应每行的索引数组
row_indices = np.arange(A.shape[0])
# 定位指定位置并执行加1操作
A[row_indices, B] += 1

说明 #

• np.arange(A.shape[0])生成与A行数一致的索引数组（此处为[0,1]），对应每行的行号
• A[row_indices, B]精准选取A中(0,2)和(1,4)这两个目标位置的元素
• 通过+=1直接对这些元素完成增量操作，一步达成需求

内容的提问来源于stack exchange，提问作者kharsair