Search code examples
pythonarraysnumpyperformancevectorization

Given two large numpy arrays, how to efficiently compute a matrix of a-b for any all (a,b) pair of the two arrays?


I have two large numpy arrays A, B. How can I efficiently (e.g. using numpy vectorization) compute a matrix C of shape len(A) * len(B) such that C[i,j]=A[i]- C[j].

For example, A is

a
b
c

B is

d
e 
f
g

Then expected C is

a-d a-e a-f a-g
b-d b-e b-f b-g
c-d c-e c-f c-g

Solution

  • import numpy as np
    
    a = np.array([1,2,3])
    b = np.array([4,5,6])
    
    A = np.tile(a.reshape(3,1),(1,3))
    B = np.tile(b,(3,1))
    C = A - B
    
    print(A)
    print(B)
    print(C)
    

    Yields:

    [[1 1 1]
     [2 2 2]
     [3 3 3]]
    [[4 5 6]
     [4 5 6]
     [4 5 6]]
    [[-3 -4 -5]
     [-2 -3 -4]
     [-1 -2 -3]]
    

    Try it online!