Search code examples
pythonarraysnumpyperformancevectorization

Cartesian product of two numpy arrays, with condition


Given 1d arrays w and x, below, I can form the Cartesian product using the following code.

import numpy as np

w = np.array([1, 2, 3, 4])
x = np.array([1, 2, 3, 4])

V1 = np.transpose([np.repeat(w, len(x)), np.tile(x, len(w))])
print(V1)

[[1 1]
 [1 2]
 [1 3]
 [1 4]
 [2 1]
 [2 2]
 [2 3]
 [2 4]
 [3 1]
 [3 2]
 [3 3]
 [3 4]
 [4 1]
 [4 2]
 [4 3]
 [4 4]]

But, I want the output, V1, to include ONLY array rows where w < x (as shown below). I can do this with loops, but I'm hoping for something a little faster.

[[1 2]
 [1 3]
 [1 4]
 [2 3]
 [2 4]
 [3 4]]

Solution

  • Approach #1

    Given ONLY array rows where w < x, which would be for pairwise combinations, here's one way to achieve the same -

    In [81]: r,c = np.nonzero(w[:,None]<x) # or np.less.outer(w,x)
    
    In [82]: np.c_[w[r], x[c]]
    Out[82]: 
    array([[1, 2],
           [1, 3],
           [1, 4],
           [2, 3],
           [2, 4],
           [3, 4]])
    

    Approach #2

    Using a purely masking based approach, it would be -

    In [93]: mask = np.less.outer(w,x)
    
    In [94]: s = (len(w), len(x))
    
    In [95]: np.c_[np.broadcast_to(w[:,None], s)[mask], np.broadcast_to(x, s)[mask]]
    Out[95]: 
    array([[1, 2],
           [1, 3],
           [1, 4],
           [2, 3],
           [2, 4],
           [3, 4]])
    

    Benchmarking

    Using comparatively larger arrays :

    In [8]: np.random.seed(0)
       ...: w = np.random.randint(0,1000,(1000))
       ...: x = np.random.randint(0,1000,(1000))
    
    In [9]: %%timeit
       ...: r,c = np.nonzero(w[:,None]<x)
       ...: np.c_[w[r], x[c]]
    11.3 ms ± 24.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    
    In [10]: %%timeit
        ...: mask = np.less.outer(w,x)
        ...: s = (len(w), len(x))
        ...: np.c_[np.broadcast_to(w[:,None], s)[mask], np.broadcast_to(x, s)[mask]]
    10.5 ms ± 275 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    
    In [11]: import itertools
    
    # @Akshay Sehgal's soln
    In [12]: %timeit [i for i in itertools.product(w,x) if i[0]<i[1]]
    105 ms ± 1.38 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)