使用NumPy加速数组区间比较

Question

我有一大堆数据点，这些数据点是在某个时间段内收集的。现在我想在另一个时间段内计算这些数据的加权平均值，方法是查看这两个时间段之间的重叠部分。我现在有一个成功的方法，下面会描述，但这个方法运行得比较慢。我希望能找到一些建议，让这个过程更快。

我有两组时间段，还有与其中一个时间段相关的数据点。我需要找出哪些来自第一个时间段的数据点落在第二个时间段内。比如，如果我的数组是这样的：

start1 = np.array([1.,6.,11.,16.,21.,26.,31.,36.])
stop1 = np.array([6.,11.,16.,21.,26,31.,36.,41.])
start2 = np.array([1.,2.,3.,4.,5.,6.,7.,8.,9.,10.,11.,12.,13.,14.,15.,16.,17.,...,39.])
stop2 = np.array([2.,3.,4.,5.,6.,7.,8.,9.,10.,11.,12.,13.,14.,15.,16.,17.,18.,...,40.])
data1 = np.array([2.,4.,3.,8.,4.,7.,2.,6.])

那么我的结果会是：

[array([ 2.]), array([ 2.]), array([ 2.]), array([ 2.]), array([ 2.,  4.]), array([ 2.,  4.]), array([ 4.]), array([ 4.]), array([ 4.]), array([ 4.,  3.]), array([ 4.,  3.]), array([ 3.]), array([ 3.]), array([ 3.]), array([ 3.,  8.]), array([ 3.,  8.]), array([ 8.]), array([ 8.]), array([ 8.]), array([ 8.,  4.]), array([ 8.,  4.]), array([ 4.]), array([ 4.]), array([ 4.]), array([ 4.,  7.]), array([ 4.,  7.]), array([ 7.]), array([ 7.]), array([ 7.]), array([ 7.,  2.]), array([ 7.,  2.]), array([ 2.]), array([ 2.]), array([ 2.]), array([ 2.,  6.]), array([ 2.,  6.]), array([ 6.]), array([ 6.]), array([ 6.])]

这是我目前使用的方法：

intervals_data = []
for i in range(0, len(start2)):
    intervals_data.append(data1[((stop1>=start2[i]) & (start1<=stop2[i]))])

我遇到的问题是，这些数组通常有大约25000个元素，所以每组数据大约需要20秒来处理。我通常有几百组数据，这样一来就要花大约一个小时才能完成。

有没有人能给我指个更快的方法？谢谢！

Answer 1

你可以试试一种基于迭代器的方法：

inter1 = zip(start1, stop1, data1)
inter2 = zip(start2, stop2)
data2 = []

curr1 = next(inter1)

for curr2 in inter2:
  vals = []
  while curr2[0] > curr1[1]:
    curr1 = next(inter1)
  vals.append(curr1[2])
  while curr2[1] > curr1[1]:
    curr1 = next(inter1)
    vals.append(curr1[2])
  data2.append(np.array(vals))

return data2

这里的重点是，你应该只在区间内向前移动，所以使用迭代器来完成这个工作就可以了。

补充说明：我这里假设你使用的是Python 3.x。如果你用的是Python 2.x，只需要把zip换成izip就可以了。

Answer 2

这个想法是使用二分查找来找到每个 start2[j] 在 stop1 中的最左边的索引 i，这样插入后 stop1 仍然是有序的，也就是说要满足 start2[j] <= stop1[i]。同样地，我们也要为每个 stop2[j] 在 start1 中找到最右边的 i，也就是要满足 stop2[j] >= stop1[i]。可以通过以下方式来实现：

start_indx = stop1.searchsorted(start2, "left")
end_indx = start1.searchsorted(stop2, "right")
result = [data1[start_indx[i]:end_indx[i]]for i in range(len(start_indx))]