CODEWITHSUNDEEP
arraysperformancepython-2.7numpy
Ananda
Ananda

Reputation: 3272

Modifying a numpy array efficiently

I have a numpy array A of size 10 with values ranging from 0-4. I want to create a new 2-D array B from this with its ith column being a vector corresponding to the ith element of A.

For example, the value 1 as the first element of A would correspond to B having a column vector [0,1,0,0,0] as it's first column. A having 4 as its third element would correspond to B having it's 3rd column as [0,0,0,1,0]

I have the following code:

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

The code is doing what it's supposed to be doing but I am sure there are more efficient ways of doing this. Can anyone please suggest some?

Upvotes: 1

Views: 70

Answers (1)

Divakar
Divakar

Reputation: 221524

That could be solved by initializing an array of zeros and then integer-indexing into it with indices from A and assigning 1s, like so -

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

Upvotes: 1

Related Questions