Search code examples
pythonnumpymatrix

NumPy: get a matrix of distances between values


If I have a set of points:

points = np.random.randint(0, 11, size=10)
print(points)

Output:

[ 5  4  9  7  4  1  2 10  4  2]

And if I want to get a matrix representing the distance from each point to each other point, I can do so like this:

def get_diff_array(values):
    dX = values - values[0]
    tmp = [dX]
    for i in range(1, len(dX)):
        tmp.append((dX - dX[i]))
    return np.array(tmp)

print(get_diff_array(points))

Output:

[[ 0 -1  4  2 -1 -4 -3  5 -1 -3]
 [ 1  0  5  3  0 -3 -2  6  0 -2]
 [-4 -5  0 -2 -5 -8 -7  1 -5 -7]
 [-2 -3  2  0 -3 -6 -5  3 -3 -5]
 [ 1  0  5  3  0 -3 -2  6  0 -2]
 [ 4  3  8  6  3  0  1  9  3  1]
 [ 3  2  7  5  2 -1  0  8  2  0]
 [-5 -6 -1 -3 -6 -9 -8  0 -6 -8]
 [ 1  0  5  3  0 -3 -2  6  0 -2]
 [ 3  2  7  5  2 -1  0  8  2  0]]

Is there a faster NumPy-specific way to calculate this? Currently it takes approximately 0.44 seconds for 10k points, which seems slow. This is really just a learning question to try to understand NumPy better.


Solution

  • You can use values[:, np.newaxis] - values[np.newaxis, :]:

    import numpy as np
    
    def diff(A):
        return A[:, np.newaxis] - A[np.newaxis, :]
    
    print(diff(np.random.randint(0, 11, size=10)))
    

    Prints

    [[  0  -7  -8  -3  -6  -6  -4  -5 -10  -3]
     [  7   0  -1   4   1   1   3   2  -3   4]
     [  8   1   0   5   2   2   4   3  -2   5]
     [  3  -4  -5   0  -3  -3  -1  -2  -7   0]
     [  6  -1  -2   3   0   0   2   1  -4   3]
     [  6  -1  -2   3   0   0   2   1  -4   3]
     [  4  -3  -4   1  -2  -2   0  -1  -6   1]
     [  5  -2  -3   2  -1  -1   1   0  -5   2]
     [ 10   3   2   7   4   4   6   5   0   7]
     [  3  -4  -5   0  -3  -3  -1  -2  -7   0]]