Более быстрый способ создания матрицы затрат

Question

Я использую Венгерский алгоритм в scipy, который принимает в качестве входных данных матрицу затрат двух наборов точек. Это просто означает, что каждый элемент в массиве x передается в функцию f с каждым элементом в массиве y. В настоящее время я реализовал это с вложенным для l oop в python. Вот базовый c пример того, что я делаю:

def f(a, b):

    return a * b

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

cost_mat = np.zeros((x.shape[0], y.shape[0]))

for i in range(x.shape[0]):
    for j in range(y.shape[0]):
        cost_mat[i, j] = f(x[i], y[j])

print(cost_mat)

>> out:
[[1., 2., 3.]
 [2., 4., 6.]
 [3., 6., 9.]]

Есть ли более быстрый способ сделать это? Например, как это векторизовать?

Answer 1

Примерно такая работа:

x = np.array([1, 2, 3], ndmin=2)
y = np.array([1, 2, 3], ndmin=2)

cost_mat = x * y.T

cost_matrix is ​​

array([[1, 2, 3],
       [2, 4, 6],
       [3, 6, 9]])

Давайте выберем оба решения с большими массивами:

x = np.random.rand(10000,1)
y = np.random.rand(10000,1)

def f(a, b):
    return a * b

# Start timing here
cost_mat1 = np.zeros((x.shape[0], y.shape[0]))

for i in range(x.shape[0]):
    for j in range(y.shape[0]):
        cost_mat1[i, j] = f(x[i], y[j])

# Wall time: 2min 13s

Использование транспонирования намного быстрее :

# Start timing here
cost_mat2 = x * y.T
# Wall time: 395 ms

А затем проверьте, что

np.array_equal(cost_mat1, cost_mat2)

Возвращает true