NumPy 中高效的三重索引矩阵计算

Question

我正在尝试在 NumPy 中运行三重矩阵索引，虽然我的代码似乎在做我想做的事，但我想知道是否有更有效的方法，因为大型矩阵的运行时间太长。

我有两个包含起点-目的地对之间旅行时间的 numpy 矩阵，第一个用于行程 1（从 i 到 k），第二个用于行程 2（从 k 到 j）。我想创建两个新矩阵：
- 第一个矩阵包含每个 i-j 对的最短旅行时间（即 i 到 k + k 到 j 的最短时间）
- 第二个矩阵显示返回此最短行程时间的最优 k 中间停靠点的索引。

我的代码如下：

nb_zones = 100
leg1 = np.random.rand(nb_zones,nb_zones)
leg2 = np.random.rand(nb_zones,nb_zones)

# initialise result matrices
total = np.zeros((nb_zones,nb_zones))
index = np.zeros((nb_zones,nb_zones))

# triple index calcs
for a in range(nb_zones):
    for b in range(nb_zones):
        # max
        total[a,b] = np.min(leg1[a] + leg2[:,b])
        # index
        index[a,b] = np.argmin(leg1[a] + leg2[:,b])

对于 800 个区域的数量，这会导致长达 30 秒的运行时间，是否有更聪明的方法来做到这一点，而无需通过所有矩阵单元进行双循环？

Answer 1

使用 numpy 广播并在 np.min 和 np.argmin 上指定一个轴可以避免一个 for 循环：

for b in range(nb_zones):
    # max
    total[:,b] = np.min(leg1 + leg2[:,b], axis=1)
    # index
    index[:,b] = np.argmin(leg1 + leg2[:,b], axis=1)

Answer 2

没有循环：

duallegs = leg1[:,:,np.newaxis] + leg2[np.newaxis, :]
index = np.argmin(duallegs, axis=1)

rowidx, colidx = np.mgrid[0:nb_zones, 0:nb_zones]
total = duallegs[(rowidx, index, colidx)]

时间：

import numpy as np

nb_zones = 800
leg1 = np.random.rand(nb_zones,nb_zones)
leg2 = np.random.rand(nb_zones,nb_zones)

def f1():
    # initialise result matrices
    total = np.zeros((nb_zones,nb_zones))
    index = np.zeros((nb_zones,nb_zones))

    # triple index calcs
    for a in range(nb_zones):
        for b in range(nb_zones):
            # max
            total[a,b] = np.min(leg1[a] + leg2[:,b])
            # index
            index[a,b] = np.argmin(leg1[a] + leg2[:,b])

    return total, index

def f2():
    duallegs = leg1[:,:,np.newaxis] + leg2[np.newaxis, :]
    index = np.argmin(duallegs, axis=1)

    rowidx, colidx = np.mgrid[0:nb_zones, 0:nb_zones]
    total = duallegs[(rowidx, index, colidx)]

    return total, index

%timeit f1()
%timeit f2()

8 s ± 819 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
4.85 s ± 360 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)