Несколько диапазонов / векторизация np.arange - PullRequest
Новая
       19

Несколько диапазонов / векторизация np.arange

4 голосов
/ 30 апреля 2019

У меня есть два массива конечных точек, которые выглядят так: 

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

Это решение медленное, и я уверен, что оно все еще может быть значительно улучшено, если полностью заменить комбо zip и list compition понимание некоторыми функциями полностью в Numpy, просто я не знаю как. Ребята, можете ли вы показать мне наиболее эффективный способ сделать это?

Заранее спасибо, ребята

Код для генерации этих двух массивов: 

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

1 Ответ

4 голосов
/ 30 апреля 2019

Вот векторизованный подход: 

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