python - 针对许多不同的开始/结束值计算排序数组中的项目数的有效方法

Question

我想计算许多不同开始/结束日期的开始日期和结束日期之间的 S&P 开放天数。

一个假例子

SPopen = pd.bdate_range(start = '1950-01-01', end = '2020-01-01')
startdates = pd.bdate_range(start = '1970-01-01', end = '2000-01-01')
enddates = startdates + pd.Timedelta(1, 'Y')

对于开始/结束日期中的每一对，我可以这样做

np.sum( (SPopen > start) & (SPopen <= end) )

获取SP开放天数，但循环几千次很慢。有没有有效的方法来做到这一点？

注意:SP 不会在所有工作日开放，并且 np.busday_count 不起作用。

Answer 1

简介

非常有趣的问题，在处理数据帧的日期时间或涉及满足间隔限制的元素计数的任何问题时可能非常有用。

建议的方法

为了解决这个问题，我们可以通过使用 np.searchsorted 来滥用数据和间隔的事实。及其可选的 'left' 和 'right' 参数。我一开始就想到了一个 NumPy 样本，这对于日期时间也有足够的概括性。

涉及的步骤

让我回顾一下我解决这个问题的历史:

1] 给定输入 -

In [618]: a  # Data array
Out[618]: array([ 0,  2,  4, 14, 15, 27, 29])

In [619]: s0  # Interval start
Out[619]: array([ 2,  6,  9, 15, 25])

In [620]: s1  # Interval stop
Out[620]: array([ 7, 10, 11, 19, 29])

2]获取左、右索引位置 -

In [621]: search_stop = np.searchsorted(a,s1,'right')
     ...: search_start = np.searchsorted(a,s0,'left')
     ...: 

3] 获取通常情况下的差异 -

In [622]: out = search_stop - search_start

4] 对于起始位置已存在于 a 中的情况，np.searchsorted(a,s0,'left') 会给我们一个较小的索引，因此它的偏移量 -

In [623]: out -= a[search_start] == s0

5] 对于情况，当间隔没有捕获任何元素时，由于最后一步偏移，我们可能会出现负计数。因此，将它们剪裁为零，这样我们就得到了所需的输出 -

In [624]: out.clip(min=0)
Out[624]: array([1, 0, 0, 0, 2])

此外，对于从数据数组中任何元素之外开始的间隔，search_start 将超出数组长度，因此请使用掩码来限制这些计算。

总结一切，我们最终会得到这样的实现 -

def vectorized_interval_count(a, s0, s1):
    search_stop = np.searchsorted(a,s1,'right')
    search_start = np.searchsorted(a,s0,'left')

    L = np.searchsorted(search_start, a.size)
    out = search_stop - search_start 
    out[:L] -= (a[search_start[:L]] == s0[:L])
    out.clip(min=0, out = out)
    return out

巨大改进

事实证明，正如 comments by OP 中提到的那样，我们可以简单地查找 'right' 索引，并且各自的差异在功能上意味着这些 left-open 和 right-close 中的元素数量> 间隔。

因此，单行解决方案是 -

np.searchsorted(a,s1,'right') - np.searchsorted(a,s0,'right')

运行时测试

在问题中给定的巨大样本数据集上进行测试，我得到了 -

In [795]: SPopen = pd.bdate_range(start = '1950-01-01', end = '2020-01-01')
     ...: startdates = pd.bdate_range(start = '1970-01-01', end = '2000-01-01')
     ...: enddates = startdates + pd.Timedelta(1, 'Y')
     ...: 

In [796]: a = SPopen
     ...: s0 = startdates
     ...: s1 = enddates
     ...: 

In [797]: out1 = [np.sum( (a > s0[i]) & (a <= s1[i]) ) for i in range(len(s0))]
     ...: out2 = vectorized_interval_count(a, s0, s1)
     ...: out3 = np.searchsorted(a,s1,'right') - np.searchsorted(a,s0,'right')
     ...: print np.allclose(out1, out2)
     ...: print np.allclose(out1, out3)
     ...: 
True
True

In [798]: %timeit [np.sum( (a > s0[i]) & (a <= s1[i]) ) for i in range(len(s0))]
1 loops, best of 3: 4.44 s per loop

In [799]: %timeit vectorized_interval_count(a, s0, s1)
1000 loops, best of 3: 842 µs per loop

In [800]: %timeit np.searchsorted(a,s1,'right') - np.searchsorted(a,s0,'right')
1000 loops, best of 3: 559 µs per loop

因此，与循环理解相比，速度接近 8,000x! (感谢OP!)