Numpy，使用另一个一维数组中表示的索引增加二维数组中的值

Question

这是我想做的一个例子：假设数组 A

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

和数组 B

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

我正在寻找一种操作，它将在数组 A 的每一行中将数组 B 索引的元素递增 1。因此结果 A 是：

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

第一行索引2增加1，第二行索引4增加1

Answer 1

您可以通过使用 numpy 中的高级索引来实现此目的：

A[np.arange(len(B)), B] += 1

这是通过使用np.arange(len(B))创建一个维度为(len(B), len(B))的2D数组来实现的，它代表行索引。 高级索引的第二个索引 B 表示列索引。 通过将 1 添加到A[np.arange(len(B)), B] ，您可以增加 B 指定的每一行中的元素。

Answer 2

在numpy中，您可以使用数组的arrange和shape来完成

import numpy as np

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

A[np.arange(A.shape[0]), B] += 1

print(A)

np.arange(A.shape[0])生成一个从0 to A.shape[0] - 1的整数数组。 A.shape[0]基本上是行

你也可以用循环来做..

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

for i, index in enumerate(B):
    A[i][index] += 1

print(A)