[Python] Fastest way to find the indices of two Numpy arrays that meet a condition

Say I have two large numpy arrays a1 and a2 (each with 10000 numbers). I want to find the indices in each array that meet the condition of f(x1, x2) > 0. To be clear, for each number in a1 (or a2), if there's any number in a2 (or a1) that meets the condition, then it is a valid number. For now I am using one loop over a1 and only using numpy vectorization for calculations on a2:

import numpy as np

a1 = np.random.rand(10000)
a2 = np.random.rand(10000)

def f(x1, x2):
return np.exp(x2 - x1) / 1.2 - 1

valid_i1_set = set()
valid_i2_set = set()
for i1 in range(len(a1)):
valid_i2 = np.nonzero(f(a1[i1], a2) > 0)[0].tolist()
valid_i2_set.update(valid_i2)
if len(valid_i2) > 0:
valid_i1_set.add(i1)
print(sorted(valid_i1_set))
print(sorted(valid_i2_set))


I imagine there should be a faster way to do this without a for loop. Making a product array of 10000x10000 numbers also doesn't feel right. Any suggestions?

Continue reading...