Question

我有两个如下的端点数组：

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理解组合来大大改进它，只是我不知道如何。你们能告诉我最有效的方法吗？

谢谢大家

生成这两个数组的代码：

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

Answer 1

这是向量化方法：

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])]

多个范围/矢量化的np.arange

1 个答案: