有效地修改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