Как добавить список с массивами двух разных размеров на основе условий

Question

Мне было интересно, как сделать это более эффективным способом для произвольных массивов, код написан на PyTorch, но это только для 1-й тензора. Спасибо!

test=[]
data=np.random.uniform(0,1,[20,])
x=torch.from_numpy(data).float()
x,_=torch.sort(x)
v=torch.rand(5).float()
v,_=torch.sort(v)
for i in range(len(x)):
    if x[i] < v[0]:
         test.append(v[0])
    elif x[i] < v[1]:
         test.append(v[1])
    elif x[i] < v[2]:
           test.append(v[2])
    elif x[i] < v[3]:
           test.append(v[3])
    else: 
        test.append(v[4])



  test

Answer 1

Вы можете использовать встроенную функцию Далее :

for i in x:
    test.append(next((e for e in v[:4] if  i < e), v[4]))

Вы также можете использовать понимание списка вместо for l oop:

s = v[:4]
d = v[4]
test = [next((e for e in s if  i < e), d)) for i in x]

если в переменной test уже есть некоторые элементы, вы можете использовать оператор i n-place +=:

test += [next((e for e in s if  i < e), d) for i in x]

Answer 2

Скромное изменение вашей функции:

def foo0(x,v):
    test = []
    for i in x:
        if i<v[0]:
            test.append(v[0])
        elif i<v[1]:
            test.append(v[1])
        elif i<v[2]:
            test.append(v[2])
        elif i<v[3]:
            test.append(v[3])
        else:
            test.append(v[4])
    return test

протестируйте его с целочисленными массивами (также отсортированными) для удобства сравнения:

In [152]: x = np.arange(20); v = np.arange(0,20,4)                                             
In [153]: x,v                                                                                  
Out[153]: 
(array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,
        17, 18, 19]), array([ 0,  4,  8, 12, 16]))
In [154]: foo0(x,v)                                                                            
Out[154]: [4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16, 16, 16, 16, 16]

In [155]: timeit foo0(x,v)                                                                     
21.2 µs ± 471 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Ваша версия повторяется на range и использование нескольких x[i] раз при 29.

A numpy версия:

def fooN(x,v):
    v1 = v.copy()
    v1[-1] = np.max(x)+1
    temp = x[:,None]<v1
    idx = np.argmax(temp, axis=1)
    return v[idx].tolist()

In [158]: fooN(x,v)                                                                            
Out[158]: [4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16, 16, 16, 16, 16]
In [159]: timeit fooN(x,v)                                                                     
27.7 µs ± 842 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Для этого небольшого образца это медленнее. Но с большей выборкой это быстрее. Однако это все еще нуждается в некоторой отладке с большими массивами. Поэтому я предлагаю это как концептуальный тест, а не окончательный ответ.

Передача списков в функцию делает ее еще быстрее

In [185]: %%timeit x1=x.tolist(); v1=v.tolist() 
     ...: foo0(x1,v1) 

5.12 µs ± 8.78 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Вот версия вашей функции, которая более компактна и общее - но немного медленнее

def foo01(x,v):
    # getting else v[4] part requires some extra effort
    test = []
    for i in x:
        done = False  
        for j in v[:-1]:
            if i<j: 
                test.append(j)
                done = True
                break
        if not done:
            test.append(v[-1])
    return test

===

для большего x (v по-прежнему 5 элементов):

In [215]: x = np.arange(2000); v = np.arange(0,2000,400)                                       
In [216]: v                                                                                    
Out[216]: array([   0,  400,  800, 1200, 1600])
In [217]: np.array(foo0(x,v))                                                                  
Out[217]: array([ 400,  400,  400, ..., 1600, 1600, 1600])
In [218]: np.array(fooN(x,v))                                                                  
Out[218]: array([ 400,  400,  400, ..., 1600, 1600, 1600])
In [219]: timeit foo0(x,v)                                                                     
1.95 ms ± 24.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [220]: timeit foo01(x,v)                                                                    
3.14 ms ± 36.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [221]: timeit fooN(x,v)                                                                     
147 µs ± 1.65 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)