如何使用索引数组访问沿着轴1的2D numpy数组元素？

yshpjwxd 于 2023-08-05 发布在其他

关注(0)|答案(2)|浏览(187)

设N是形状为(m,n)的numpy数组，A是索引(m,)的数组，包含0和n-1之间的值，包括。
如果I i是从0到m-1包括的索引，那么我想要访问元素N[i,A[i]]。
举例来说：

import numpy as np
m = 2000
n = 10
N = np.zeros((m,n))
A = np.random.choice(n, m)
N[???, A] = 1

字符串
我期望上面的示例代码生成一个形状为(m,n)的数组N，最初所有的零，然后设置N[i,j]=1，其中A[i]=j。
我尝试了N[:, A]=1，但它将数组N的所有元素都更改为1，这不是我想要实现的。

2条答案

使用:索引所有行，所以执行N[:,A] = 1是告诉numpy获取所有行，并在A列中将值设置为1。您只需要一对一的行和列选择。最简单的方法是使用np.arange。

N[np.arange(m), A] = 1

字符串
现在，A中的每个值都有一个对应的行，这导致m行和列选择。

您可以按如下方式设置条件N[i,j]=1 where A[i] == j：

# Using conditional selection 
condition = np.arange(m) == A
N[condition] = 1

字符串