简单的Numpy矢量化

Question

我有两个1D numpy的阵列start和stop,两者都包含整数(这是用来索引一些其它阵列).我有以下代码.

index_list = []
for i in range(len(start)):
    temp = range(start[i], stop[i])
    index_list.extend(temp)
index_list = np.array(index_list)

是否有一种简单的方法可以对此进行矢量化？

Answer 1

您可以按如下方式对其进行矢量化:

def make_index_list(start, stop):
    lens = stop - start
    cum_lens = np.cumsum(lens)
    # Sequential indices the same length as the expected output
    out = np.arange(cum_lens[-1])
    # Starting index for each section of `out`
    cum_lens = np.concatenate(([0], cum_lens[:-1]))
    # How much each section of `out` is off from the correct value
    deltas = start - out[cum_lens]
    # Apply the correction
    out += np.repeat(deltas, lens)

    return out

有一些组成的数据:

start = np.random.randint(100, size=(100000,))
stop = start + np.random.randint(1, 10 ,size=start.shape)

我们可以将代码用于测试:

In [39]: %%timeit
   ....: index_list = []
   ....: for i in range(len(start)):
   ....:     temp = range(start[i], stop[i])
   ....:     index_list.extend(temp)
   ....: index_list = np.array(index_list)
   ....:
10 loops, best of 3: 137 ms per loop

In [40]: %timeit make_index_list(start, stop)
100 loops, best of 3: 9.27 ms per loop

In [41]: np.array_equal(make_index_list(start, stop), index_list)
Out[41]: True

所以它是正确的,大约快15倍,一点都不差......