有效地修改numpy数组

Question

我有一个大小为10的numpy数组A，值范围从0-4。我想创建一个新的二维数组B，其第i列是一个对应于A的第i个元素的向量。

例如，作为A的第一个元素的值1将对应于具有列矢量[0,1,0,0,0]，因为它是第一列。具有4作为其第三个元素将对应于B有它的第三列[0,0,0,1,0]

我有以下代码：

import numpy as np 

A = np.random.randint(0,5,10) 
B = np.ones((5,10)) 

iden = np.identity(5, dtype=np.float64) 

for i in range(0,10): 
    a = A[i] 
    B[:,i:i+1] = iden[:,a:a+1] 

print A 
print B 

的代码是做什么它应该做的，但我相信有更有效的方式来做到这一点。任何人都可以请建议一些？

Answer 1

可能由从A初始化为零的数组，然后整数索引到它的索引和分配1s，像这样来解决 -

M,N = 5,10 
A = np.random.randint(0,M,N) 
B = np.zeros((M,N)) 
B[A,np.arange(len(A))] = 1