lara_toff
lara_toff

Reputation: 442

Using numpy where to return an element in the same row but different column

I am trying to use np.where to swap one element for another, in the same row (but different column).

Here is what I have so far:

arr = np.random.random_integers(0,10,size=(10,10))
split = np.random.random_integers(0,10,size=arr[1,:].shape)
newloc = np.random.random_integers(0,10,size=arr[1,:].shape)
arr2 = np.where(arr>split,arr[0,newloc],arr)

My problem is that 0 in arr[0,newloc] means that it always pulls from row 0 off arr. But I want something like "same_row" if that makes sense.

Upvotes: 0

Views: 238

Answers (1)

Ivan
Ivan

Reputation: 40708

I think what you're looking for is arr[:, newloc]. This will allow you to index arr's second axis (i.e. the columns) with newloc's values.

np.where(arr > split, arr[:, newloc], arr)

Here is an example with two rows:

>> arr = np.random.randint(0, 10, size=(2,10))
>> split = np.random.randint(0, 10, size=arr.shape[1])
>> newloc = np.random.randint(0, 10, size=arr.shape[1])

>> arr
array([[2, 3, 0, 0, 1, 4, 6, 5, 1, 9],
       [9, 2, 6, 3, 2, 6, 3, 6, 2, 5]])

>> split
array([5, 7, 3, 5, 1, 6, 0, 8, 8, 6])

>> newloc
array([1, 7, 0, 5, 0, 7, 3, 6, 4, 8])

>> np.where(arr > split, 1, 0) # mask
array([[0, 0, 0, 0, 0, 0, 1, 0, 0, 1],
       [1, 0, 1, 0, 1, 0, 1, 0, 0, 0]])

>> np.where(arr > split, arr[:, newloc], -1)
array([[-1, -1, -1, -1, -1, -1,  0, -1, -1,  1],
       [ 2, -1,  9, -1,  9, -1,  3, -1, -1, -1]])

There, you can check if the columns are being replaced correctly.


I have used np.random.randint in the above snippet since np.random.random_integers is deprecated.

Upvotes: 1

Related Questions