Home > front end >  Numpy compare with 2 different dimension array
Numpy compare with 2 different dimension array

Time:04-22

a = np.array([[  0, 100,   0], 
       [  0,   0,   0], 
       [  0,  50,   0]])
b = np.array([0, 50, 0])

c = np.array([0, 0, 1])

How can I get array c through a and b, except use for? If b equals the item in a, then the same index item of c should be 1.

This question bothers me a lot when I compare a pixel in an image. When I use for statement, it will be too slow.

CodePudding user response:

Here's a possible way. First, b is subtracted from each row of a and the absolute value at each index is found. Then, the sum of each row is taken, and if the sum is not 0, then that value in c becomes 0. If the sum is 0, then that index row in a is equal to b and the index in c becomes 1.

The absolute value step is necessary in case you have a row in a like [b[0] - 4, b[1] 4, 0], because the sum of that row minus b would actually be 0 because of the 4 and -4, even though that row and b aren't the same. That's why you need the absolute value step.

c = np.sum(np.abs(a - b), axis=1)
c = np.where(c != 0, 0, 1)

I tested for a with shape (100000, 3) and b with shape (3,). For my method, the time taken is 0.0085 seconds, while a for loop takes 0.0408 seconds, so more than a 4x speedup.

Note that for small a (less than ~300 rows, always with 3 columns), the for loop is faster.

CodePudding user response:

It looks to me like you want to output True for a row in a if all values in that row are equal to b?

You can do that like this:

result = (a == b).all(axis=1)

result will be

[False False True]
  • Related