CODEWITHSUNDEEP
pythonnumpy
Shan
Shan

Reputation: 19243

comparing ndarray with values in 1D array to get a mask

I have two numpy array, 2D and 1D respectively. I want to obtain a 2D binary mask where each element of the mask is true if it matches any of the element of 1D array.

Example

 2D array 
 -----------
 1 2 3 
 4 9 6
 7 2 3


 1D array 
 -----------
 1,9,3

 Expected output
 ---------------
 True   False  True
 False  True   False
 False  False  True

Thanks

Upvotes: 2

Views: 1110

Answers (2)

unutbu
unutbu

Reputation: 879691

You could use np.in1d. Although np.in1d returns a 1D array, you could simply reshape the result afterwards:

In [174]: arr = np.array([[1,2,3],[4,9,6],[7,2,3]])

In [175]: bag = [1,9,3]

In [177]: np.in1d(arr, bag).reshape(arr.shape)
Out[177]: 
array([[ True, False,  True],
       [False,  True, False],
       [False, False,  True]], dtype=bool)

Note that in1d is checking of the elements in arr match any of the elements in bag. In contrast, arr == bag tests if the elements of arr equal the broadcasted elements of bag element-wise. You can see the difference by permuting bag:

In [179]: arr == np.array([1,3,9])
Out[179]: 
array([[ True, False, False],
       [False, False, False],
       [False, False, False]], dtype=bool)

In [180]: np.in1d(arr, [1,3,9]).reshape(arr.shape)
Out[180]: 
array([[ True, False,  True],
       [False,  True, False],
       [False, False,  True]], dtype=bool)

When you compare two arrays of unequal shape, NumPy tries to broadcast the two arrays to a single compatible shape before testing for equality. In this case, [1, 3, 9] gets broadcasted to 

array([[1, 3, 9],
       [1, 3, 9],
       [1, 3, 9]])

since new axes are added on the left. You can check the effect of broadcasting this way:

In [181]: np.broadcast_arrays(arr, [1,3,9])
Out[185]: 
[array([[1, 2, 3],
        [4, 9, 6],
        [7, 2, 3]]), 
 array([[1, 3, 9],
        [1, 3, 9],
        [1, 3, 9]])]

Once the two arrays are broadcasted up to a common shape, equality is tested element-wise, which means the values in corresponding locations are tested for equality. In the top row, for example, the equality tests are 1 == 1, 2 == 3, 3 == 9. Hence,

In [179]: arr == np.array([1,3,9])
Out[179]: 
array([[ True, False, False],
       [False, False, False],
       [False, False, False]], dtype=bool)

Upvotes: 2

Julien Spronck
Julien Spronck

Reputation: 15433

a = np.array([[1,2,3],[4,9,6],[7,2,3]])
b = np.array([1,9,3])

Have you tried this:

print a == b
## array([[ True, False,  True],
##        [False,  True, False],
##        [False, False,  True]], dtype=bool)

Look up broadcasting (http://docs.scipy.org/doc/numpy-1.10.1/user/basics.broadcasting.html) to see why this works.

Upvotes: 1

Related Questions