将一维整数数组与二维数组中的行进行比较

Question

I have a piece of code that is running, but is currently a bottleneck, and I was wondering whether there is a smarter way to do it.

I have a 1D array of integers between 0-20, with a length of 20-1000 (it varies) and I'm trying to compare it to a set of 1D arrays that are stored in a 2D array. I wish to find any row in the 2D array that completely matches the 1D array.

My current approach to do this is the following:

res = np.mean(one_d_array == two_d_array,axis=1) == 1

The problem with this approach is that it will compare all elements in all rows, even if these rows don't match on the first element, second, ect... Which is of course very inefficient. I could remedy this by looping through the rows and comparing each row individually, then I would probably be able to stop the comparison as soon as one element is false. However then I would be stuck with a slow for loop, which would also not be ideal.

So I'm wondering is there some other clever way to get the best of both of these approaches?

Answer 1

numpy has a few useful built-in functions for checking matrix/vector equality, this is about twice as fast:

import numpy as np
import time
x = np.random.random((1, 1000))
y = np.random.random((10000, 1000))
y[53] = x

t = time.time()
x_in_y = np.equal(x, y).all(axis=1)  # equal(x, y) returns a row x col matrix of True for matches; all(axis=0) returns a vector len(rows) if the entire row in x == y is true
idx = np.where(x_in_y)  # returns the indicies where x_in_y is true (here 53)
print(time.time() - t)  # 0.019975900650024414

t = time.time()
res = np.mean(x == y, axis=1) == 1
print(time.time() - t)  # 0.03999614715576172