CODEWITHSUNDEEP
pythonarrayspython-2.7list-comprehension
Nev Pires
Nev Pires

Reputation: 175

replace elements in a m x n array using list comprehension

I have an array: 

a = array([[1,2,3,1], [2,5,3,1], [0,0,0,0], [5,3,2,5]])

I want to iterate through the array, based off the last item in each row (0 or 1). If the last item in each row is 0, I want to change the 3 other items (only) in the row to np.nan.

For example: 

a = array([[1,2,3,1], [2,5,3,1], [nan, nan, nan, 0], [5,3,2,5]])

I can do this using a for loop. i.e.: 

for frames in range(len(a)):
    if a[frames][3] == 0:
        a[frames][0:2] = np.nan

Is there a more efficient way to do this using list comprehension? So far this is all I've come up with, but feel that it could be far more efficient: 

a = np.array([[np.nan, np.nan, np.nan] if frames[3] == 0 else frames[0:3] for frames in a])

As this will create an array of arrays, as well as crop the last column

Thanks in advance!

Upvotes: 3

Views: 680

Answers (1)

rafaelc
rafaelc

Reputation: 59274

IIUC, you don't need a list comprehension. Use indexing

>>> a = np.array([[1,2,3,1], [2,5,3,1], [0,0,0,0], [5,3,2,5]], dtype=float)
>>> a[a[:,-1] == 0, 0:3] = np.nan

array([[  1.,   2.,   3.,   1.],
       [  2.,   5.,   3.,   1.],
       [ nan,  nan,  nan,   0.],
       [  5.,   3.,   2.,   5.]])

If you have a dict of these arrays, just index each of them

for a in data.values():
    a[a[:,-1] == 0, 0:3] = np.nan

Upvotes: 2

Related Questions