CODEWITHSUNDEEP
pythonarraysnumpy
Shew
Shew

Reputation: 1596

Numpy assign an array value based on the values of another array with column selected based on a vector

I have a 2 dimensional array

X
array([[2, 3, 3, 3],
       [3, 2, 1, 3],
       [2, 3, 1, 2],
       [2, 2, 3, 1]])

and a 1 dimensional array 

y
array([1, 0, 0, 1])

For each row of X, i want to find the column index where X has the lowest value and y has a value of 1, and set the corresponding row column pair in a third matrix to 1

For example, in case of the first row of X, the column index corresponding to the minimum X value (for the first row only) and y = 1 is 0, then I want Z[0,0] = 1 and all other Z[0,i] = 0. Similarly, for the second row, column index 0 or 3 gives the lowest X value with y = 1. Then i want either Z[1,0] or Z[1,3] = 1 (preferably Z[1,0] = 1 and all other Z[1,i] = 0, since 0 column is the first occurance)

My final Z array will look like

Z
array([[1, 0, 0, 0],
       [1, 0, 0, 0],
       [1, 0, 0, 0],
       [0, 0, 0, 1]])

Upvotes: 1

Views: 1031

Answers (2)

Paritosh Singh
Paritosh Singh

Reputation: 6246

One way to do this is using masked arrays.

import numpy as np

X = np.array([[2, 3, 3, 3],
              [3, 2, 1, 3],
              [2, 3, 1, 2],
              [2, 2, 3, 1]])

y = np.array([1, 0, 0, 1])
#get a mask in the shape of X. (True for places to ignore.)
y_mask = np.vstack([y == 0] * len(X))

X_masked = np.ma.masked_array(X, y_mask)

out = np.zeros_like(X)

mins = np.argmin(X_masked, axis=0)
#Output: array([0, 0, 0, 3], dtype=int64)

#Now just set the indexes to 1 on the minimum for each axis.
out[np.arange(len(out)), mins] = 1

print(out)
[[1 0 0 0]
 [1 0 0 0]
 [1 0 0 0]
 [0 0 0 1]]

Upvotes: 2

Zachi Shtain
Zachi Shtain

Reputation: 836

you can use numpy.argmin(), to get the indexes of the min value at each row of X. For example:

import numpy as np
a = np.arange(6).reshape(2,3) + 10
ids = np.argmin(a, axis=1)

Similarly, you can the indexes where y is 1 by either numpy.nonzero or numpy.where. Once you have the two index arrays setting the values in third array should be quite easy.

Upvotes: 0

Related Questions