CODEWITHSUNDEEP
pythonarraysloopsnumpy
Rose
Rose

Reputation: 203

How to create new array based on position in python?

I have an array (b) and I want to add new rows according to the position of 1s in the (a) array.

a=np.array([1,1,0,1,0]

b=np.array([1,2,3,4,5])

I'd need to make a new array like this:

Output: array([1,2,3,4,5],[1,0,0,0,0],[0,1,0,0,0],[0,0,0,1,0])

How could this be done? Loop and append? Many thanks in advance

Upvotes: 0

Views: 704

Answers (2)

Wei Hu
Wei Hu

Reputation: 60

def solution(a, b):
    result = [b]
    for i, elem in enumerate(a):
        if elem == 1:
            temp = [0] * len(a)
            temp[i] = 1
            result.append(temp)
    return np.array(result)

Upvotes: 1

Divakar
Divakar

Reputation: 221574

One-liner with broadcasting -

In [14]: np.vstack((b,np.arange(len(a)) == np.flatnonzero(a)[:,None]))
Out[14]: 
array([[1, 2, 3, 4, 5],
       [1, 0, 0, 0, 0],
       [0, 1, 0, 0, 0],
       [0, 0, 0, 1, 0]])

With zeros-initialized array -

In [17]: s = a.sum()

In [18]: c = np.zeros((s,len(a)),dtype=b.dtype)

In [20]: c[np.arange(s),np.flatnonzero(a)] = 1

In [21]: np.vstack((b,c))
Out[21]: 
array([[1, 2, 3, 4, 5],
       [1, 0, 0, 0, 0],
       [0, 1, 0, 0, 0],
       [0, 0, 0, 1, 0]])

Upvotes: 1

Related Questions