如何同时使用数组索引和切片索引NumPy数组

Question

我有一个包含索引的数组和一个包含值的数组。

对于索引数组中的每一个索引，我需要高效地在numpy数组中生成一行结果。

结果数组中的这一行需要包含索引在值数组中的前一个值、当前索引对应的值，以及接下来的两个值。

这是我用循环实现的方式：

import numpy as np

indexes = np.array([5,10,12])
values =  np.array([ 0.89643977,  0.50794841,  0.75995795,  0.78029348,  0.83609961,
        0.47534985,  0.17330516,  0.15152753,  0.15504392,  0.10245308,
        0.70428183,  0.36804107,  0.13074141,  0.77377332,  0.11368238,
        0.74711745,  0.89431082,  0.35544423,  0.08592396,  0.28762   ])

result = values[indexes[0] -1:index[0] -1 + 4]
for index in indexes[1:]:
    result = numpy.vstack([result, values[index -1:index -1 + 4]])

print result
[[ 0.83609961  0.47534985  0.17330516  0.15152753]
 [ 0.10245308  0.70428183  0.36804107  0.13074141]
 [ 0.36804107  0.13074141  0.77377332  0.11368238]]

我的索引和数值数组都很大，循环的执行时间太长了。有没有办法不使用循环来实现这个？

Answer 1

你为什么选择 values[index -1:][:4] 而不是 values[index -1: index -1 + 4] 呢？我猜你可能更想把所有结果放到一个列表里，然后再用 np.vstack 处理，而不是在循环中一个一个处理。等有更好的答案之前，你可以考虑使用：

ls_results = [values[index -1:index -1 + 4] for index in indexes]
result = np.vstack(ls_results)

Answer 2

只要没有索引超出范围，这样做就可以正常工作，而且应该会快很多：

In [4]: idx = indexes[:, None] + np.arange(-1, 3)

In [5]: values[idx]
Out[5]:
array([[ 0.83609961,  0.47534985,  0.17330516,  0.15152753],
       [ 0.10245308,  0.70428183,  0.36804107,  0.13074141],
       [ 0.36804107,  0.13074141,  0.77377332,  0.11368238]])