有效地修改numpy数组

时间:2017-05-05 14:29:58

标签: arrays performance python-2.7 numpy

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

例如,作为A的第一个元素的值1将对应于具有列向量[0,1,0,0,0]的B作为其第一列。将4作为其第三个元素对应于B将其第3列作为[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

代码正在做它应该做的事情,但我相信有更有效的方法可以做到这一点。任何人都可以建议一些吗?

1 个答案:

答案 0 :(得分: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