Ускорение обработки списка в Python

Question

У меня есть два списка list1 и list2, каждый размером 5000, причем каждая запись в списках является numpy.array. Я хочу быстро и эффективно вычислить квадраты евклидовых расстояний между элементами списков, т.е. мне нужно вычислить sum((list1[i]-list2[j])**2) для каждой комбинации i и j, что в общей сложности составляет 2500000 комбинаций. В настоящее время я сделал это, запустив двойной l oop и записав каждый результат в 2d numpy.array с помощью

result[i,j] = sum((list1[i]-list2[j])**2) 

, но мне все еще нужно на моем компьютере около 4 минут времени. Мне было интересно, можно ли использовать какие-нибудь уловки для дальнейшего ускорения вычислений.

Answer 1

Если вы настаиваете на numpy (при условии, что ваши внутренние массивы 1-D):

dist_mat = ((list1[:,None,:]-list2[:,:])**2).sum(2)

Обратите внимание, что согласно вашему определению расстояния, о котором идет речь, это квадрат евклидовых расстояний. Если вам нужно само расстояние, просто возьмите квадрат root этого.
В противном случае я бы предпочел комментарий @ Quang:

from scipy.spatial import distance_matrix
dist_mat = distance_matrix(list1, list2)