2017-01-03 1 views
1

У меня есть два массива одинаковой формы:Python скалярное произведение с караульной условием

a = numpy.array([7, 5, 0, 2, 9, 6, 4, 1]) 
b = numpy.array([1, 6, 3, 9, 1, 1, 3, 8]) 

я ищу хороший способ создания суммы продукта a и b, но такое, что вычисление останавливается перед достигая предела l, а затем возвращает элементы в a, которые использовались для этого расчета.

Таким образом, в приведенном выше примере, если я исполню:

numpy.dot(a,b)

я 90. Тем не менее, в качестве примера, если l=50 то, что я хотел бы вместо того, чтобы функция, которая ведет себя следующим образом:

>>> foo(a,b,l) 
[7,5] 
+0

Определение "хороший" - Short/производительным? – Divakar

+0

ммм. Хороший вопрос ... если бы мне пришлось выбирать, исполнитель. спасибо –

ответ

2

Один подход с cumsum получить сложений поэлементного умножений, то argmax получить первый индекс предельного суммирования продукта и, наконец, slicing в a -

a[:((a*b).cumsum()>l).argmax()] 

пробег образца -

# Adding 2 as the third elem in 'a' for a better test 
In [103]: a = np.array([7, 5, 2, 2, 9, 6, 4, 1]) 
    ...: b = np.array([1, 6, 3, 9, 1, 1, 3, 8]) 
    ...: 

In [104]: l = 50 

In [105]: (a*b).cumsum() 
Out[105]: array([ 7, 37, 43, 61, 70, 76, 88, 96]) 

In [106]: a[:((a*b).cumsum()>l).argmax()] 
Out[106]: array([7, 5, 2]) 

Продолжительность испытания -

In [116]: a = np.random.randint(0,4,(10000)) 

In [117]: b = np.random.randint(0,4,(10000)) 

In [118]: l = 10000 

In [119]: %timeit a[np.where((a*b).cumsum() < l)] # @wildwilhelm's soln 
10000 loops, best of 3: 92.8 µs per loop 

In [121]: %timeit a[(a*b).cumsum() < l] # @wildwilhelm's soln w/o np.where 
10000 loops, best of 3: 82.6 µs per loop 

In [122]: %timeit a[:((a*b).cumsum()>l).argmax()] 
10000 loops, best of 3: 71.3 µs per loop 
2
>>> import numpy as np 
>>> l = 50 
>>> a[np.where((a*b).cumsum() < l)] 
array([7, 5, 0]) 
+0

Вы можете избежать этого 'np.where'. Добавлено, что в тесте runtime в моем посте показано некоторое улучшение. – Divakar

+0

@ Дикавар +1 для выбора времени:) – wildwilhelm

Смежные вопросы