CODEWITHSUNDEEP
pythonarraysnumpyindexingnumpy-ndarray
zerohedge
zerohedge

Reputation: 3765

numpy: convert multiple assignments to a single one using OR

taxi_modified is a two-dimensional ndarray.

Code below works, but seems un-pythonic:

taxi_modified[taxi_modified[:, 5] == 2, 15] = 1
taxi_modified[taxi_modified[:, 5] == 3, 15] = 1
taxi_modified[taxi_modified[:, 5] == 5, 15] = 1

Need to assign 1 to col at index 15 if col at index 5 is 2, 3, or 5.

The below didn't work:

taxi_modified[taxi_modified[:, 5] == 2 | 3 | 5, 15] = 1

Upvotes: 1

Views: 37

Answers (1)

jpp
jpp

Reputation: 164773

You can use fancy indexing with np.isin (NumPy v1.13+), or np.in1d for older versions.

Here's a demo:

# example input array
A = np.arange(16).reshape((4, 4))

# calculate Boolean mask for rows
mask = np.isin(A[:, 1], [1, 5, 13])

# assign values, converting mask to integers
A[np.where(mask), 2] = -1

print(A)

array([[ 0,  1, -1,  3],
       [ 4,  5, -1,  7],
       [ 8,  9, 10, 11],
       [12, 13, -1, 15]])

In one line, this can be written:

A[np.where(np.isin(A[:, 1], [1, 5, 13])), 2] = -1

Upvotes: 2

Related Questions