CODEWITHSUNDEEP
pythonarraysnumpy
user3357979
user3357979

Reputation: 627

Create a new numpy array based on conditions set out in numpy array

How can I create an array based on certain conditions in another array. For example, if I have an array that gives me a Base number, a start and end number, and then multiple other Base numbers. I want to create a new matrix that lists the Base number, the loop number (based on start/end) and then the other Base number associated with this, while ignoring 0's. I am trying to find a way to do this without using a for loop.

For example, how can I get array B from array A.

           Base Start End Base1 Base2 Base3
A=np.array([[100,  1,   2,  101,  102,  103],
            [101,  3,   4,  100,  103,    0]])


B=np.array[[100,1,101,1],
           [100,1,102,1],
           [100,1,103,1],
           [100,2,101,2],
           [100,2,102,2],
           [100,2,103,2],
           [101,3,100,3],
           [101,3,103,3],
           [101,4,100,4],
           [101,4,103,4]]

Thanks for the help!

Upvotes: 4

Views: 948

Answers (2)

Eelco Hoogendoorn
Eelco Hoogendoorn

Reputation: 10769

Here you go: pure numpy, no for loops at the python level

C = np.empty((len(A), 2,3,4), np.int)
C[...,0] = A[:,0  ][:,None,None]
C[...,1] = A[:,1:3][:,:   ,None]
C[...,2] = A[:,3: ][:,None,:   ]
C[...,3] = np.arange(len(A)*2).reshape(len(A),2,1)+1

C = C.reshape(-1,4)
C = C[C[:,2]!=0]

print C

edit: minor cleanup

edit2: or for maximum brevity and crypticity (see comment):

C = np.empty((len(A), 2,3,2,2), np.int)
C[...,0,0] = A[:,0  ][:,None,None]
C[...,1,0] = A[:,3: ][:,None,:   ]
C[...,:,1] = A[:,1:3][:,:   ,None,None]

Upvotes: 2

ysakamoto
ysakamoto

Reputation: 2532

Here you go... nasty list comprehension

A = array([[100,   1,   2, 101, 102, 103], [101,   3,   4, 100, 103,   0]])
B = [[A[i,0], b, c, b] for i in range(len(A)) for b in A[i,1:3] for c in A[i,3:6] if c != 0]

>>> B
[[100, 1, 101, 1],
 [100, 1, 102, 1],
 [100, 1, 103, 1],
 [100, 2, 101, 2],
 [100, 2, 102, 2],
 [100, 2, 103, 2],
 [101, 3, 100, 3],
 [101, 3, 103, 3],
 [101, 4, 100, 4],
 [101, 4, 103, 4]]

Upvotes: 2

Related Questions