CODEWITHSUNDEEP
pythonarraysnumpyoptimization
Alexandru Dinu
Alexandru Dinu

Reputation: 1267

Possibly faster code using a numpy construct for comparison of two arrays

I have the following code which calculates the number of pairs 1 <= i < j <= n such that xs[i] == ys[j]:

def f(xs, ys):
    s = 0
    for j in range(xs.size):
        s += np.sum(xs[:j] == ys[j])
    return s

This is called a couple of times from some other procedure, so I need it to be as fast as possible (at some memory cost).

The sizes of the arrays are > 1e6.

Is there a faster equivalent that makes use of some numpy magic which gets rid of the for loop?

Upvotes: 1

Views: 76

Answers (1)

Ehsan
Ehsan

Reputation: 12407

One way of doing it if xs and ys are of the same size:

s = np.triu(xs[:,None]==ys,1).sum()

And if xs and ys are different sizes (according to your code, you only compare same length of ys with xs. In case you want to compare xs with ALL ys, use line above):

s = np.triu((xs[:,None]==ys[:xs.size]),1).sum()

or equivalently:

s = (xs[:,None]==ys)[np.triu_indices(xs.size,1)].sum()

Which compares all elements of a 2-D xs with ys and sum the equals on upper triangle (same as loop's inner line)

If your arrays are too large and you run into memory issue, simply chunk your arrays and use lines above to compare blocks on diagonal plus all of the off-diagonal blocks on upper triangle and add them up.

Upvotes: 1

Related Questions