多个范围/矢量化np.arange

Posted 2021-04-09

我有两个端点数组，如下所示：

t1 = np.array([0,13,22,...,99994])
t2 = np.array([4,14,25,...,99998])

我正在寻找最有效的方法来生成如下所示的输出：

np.array([0,1,2,3,4,13,14,22,23,24,25,...,99994,99995,99996,99997,99998])

一种方法是这样的：

np.array([i for a, b in zip(t1, t2) for i in range(a, b + 1)])

这个解决方案很慢，我确信它仍然可以通过完全用Numpy完全替换zip和list comprehension组合来完全改进，这只是我不知道如何。你们能告诉我最有效的方法吗？

提前谢谢你们

---

用于生成这两个数组的代码：

import numpy as np

m =10000
Z = np.arange(0,10*m,10)

t1 = np.random.randint(5, size =m ) + Z
t2 =np.random.randint(5,size = m) + 5 + Z

答案

这是一个矢量化方法：

---

def n_ranges(start, end, return_flat=True):
    '''        
    Returns n ranges, n being the length of start (or end,
    they must be the same length) where each value in
    start represents the start of a range, and a value 
    in end at the same index the end of it
    ----
    a: np.array
       1D array representing the start of a range. 
       Each value in start must be <= than that
       of stop in the same index
    ----       
    Returns:
      All ranges flattened in a 1darray if return_flat is True
      otherwise an array of arrays with a range in each
    '''
    # lengths of the ranges
    lens = end - start
    # repeats starts as many times as lens 
    start_rep = np.repeat(start, lens)
    # helper mask with as many True in each row 
    # as value in same index in lens
    arr = np.arange(lens.max())
    m =  arr < lens[:,None]
    # ranges in a flattened 1d array
    # right term is a cumcount up to each len
    ranges = start_rep + (arr * m)[m]
    # returns if True otherwise in split arrays
    if return_flat:
        return ranges
    else:
        return np.split(ranges, np.cumsum(lens)[:-1])

---

样品运行：

t1 = np.array([0,13,22])
t2 = np.array([4,14,25])
n_ranges(t1, t2+1)
# array([ 0,  1,  2,  3,  4, 13, 14, 22, 23, 24, 25], dtype=int32)

并设置return_flat = False：

n_ranges(t1, t2+1, return_flat=False)
# [array([0, 1, 2, 3, 4]), array([13, 14]), array([22, 23, 24, 25])]