NumPy - fastest lazy lexicographical comparing of 1D arrays

Viewed 210

I have two NumPy 1D arrays a and b.

How do I compare them lexicographically? Meaning that 1D arrays should be compared same way as Python compares tuples.

Main thing is that this should be done lazily, i.e. function should return result as soon as it is found on the left-most occurance of known result.

Also I'm looking for the fastest solution for numpy arrays. For some vectorized implementation, maybe using other numpy functions.

Otherwise non-lazy simple implementation could be like this:

i = np.flatnonzero((a < b) != (a > b))
print('a ' + ('==' if i.size == 0 else '<' if a[i[0]] < b[i[0]] else '>') + ' b')

Or lazy simple variant but slow due to using pure Python types:

ta, tb = tuple(a), tuple(b)
print('a ' + ('<' if ta < tb else '==' if ta == tb else '>') + ' b')

Another solution would be to use np.lexsort, but the question is if it is optimized for the case of just two columns (two 1D arrays) or not, also if it is lazy at all? Also the question is that lexsort's result is probably not enough to have three possibilities of answer </==/>, probably it is only enough to tell if <=. Also lexsort needs some non-lazy preprocessing like np.stack and reversing rows order.

print('a ' + ('<=' if np.lexsort(np.stack((a, b), 1)[::-1])[0] == 0 else '>') + ' b')

But can it be implemented in numpy lazily and fast? I need lazy behavior because 1D arrays can be quite large but in most cases comparison result is known very close to the beginning.

2 Answers

In straight python you'd iterate over the zipped lists:

def lazy_compare(a, b):
    for x, y in zip(a, b):
        if x < y:
            return 'a < b'
        if x > y:
            return 'a > b'
    return 'a == b'

e.g.

print(lazy_compare(['a', 'b', 'c', 'd', 'e'], ['a', 'b', 'b', 'd', 'e']))
print(lazy_compare(['a', 'b', 'c', 'd', 'e'], ['a', 'b', 'c', 'd', 'f']))
print(lazy_compare(['a', 'b', 'c', 'd', 'e'], ['a', 'b', 'c', 'd', 'e']))

Output:

a > b
a < b
a == b

Since zip returns an iterator that only generates the values as you use them, this is lazy and will return a result as soon as it finds one, so will only require going over the entirety of both lists if they are equal.

One might guess that using a loop and indexing the arrays could plausibly be faster than zip, but this is not the case.

Take these definitions for the basis of comparison.

def lex_leq_zip(a, b):
    for x, y in zip(a, b):
        if x > y:
            return False
    return True

def lex_leq_index(x,y):
    for i in np.arange(x.size):
        if x[i] > y[i]:
            return False
    return True

And then we scan through different sizes of arrays to collect data on changes:

for L in range(1,100000, 1000):
    for rep in range(10):
        x = np.random.random(size=L)
        y = np.random.random(size=L)
        z = timeit('lex_leq_zip(x,y)',
              globals={'lex_leq_zip':lex_leq_zip,
                       'x':x,
                       'y':y},
              number=1)
        i = timeit('lex_leq_index(x,y)',
              globals={'lex_leq_index':lex_leq_index,
                       'x':x,
                       'y':y},
              number=1)
        plt.scatter([L], [z], color='k')
        plt.scatter([L], [i], color='b')
plt.show()

Zooming in on the resulting plot, I obtained this: enter image description here

Recalling from the code above that the vertical axis is time in seconds, the horizontal axis is the length of the array, the blue factor is the index-based implementation, and the black factor is the zip-based implementation. While we're considering pretty small fractions of a second (which can be precious in some contexts), it is pretty clear that the zip-based approach is faster.

Footnote: I also tried using Numba's @jit(nopython=True) decorator on the index-based implementation, however it showed a similar pattern.

Footnote: I also tried NumPy's np.vectorize on both implementations, but actually both result in errors pertaining to trying to index a number.

Related