Алгоритмы Numpy векторизации для суммирования чисел с одинаковыми отметками времени

Question

У меня есть два массива P и T. P [i] - число, отметка времени которого T [i]; Там могут быть дублированные метки времени.

Я хочу создать еще два массива Q и U, где Q [i] имеет метку времени U [i], а Q [i] - сумма всех элементов в P, имеющих метку времени U [i];

Например, для

P = [1, 2, 3, 4, 5] T = [0, 0, 1, 1, 1]

буду производить

Q = [3, 12] U = [0, 1];

Есть ли быстрый способ сделать это в numpy, который, мы надеемся, векторизует это?

Answer 1

Использование numpy 1.4 или лучше:

import numpy as np

P = np.array([1, 2, 3, 4, 5]) 
T = np.array([0, 0, 1, 1, 1])

U,inverse = np.unique(T,return_inverse=True)
Q = np.bincount(inverse,weights=P)
print (Q, U)
# (array([  3.,  12.]), array([0, 1]))

---

Обратите внимание, что это , а не самое быстрое решение. Я проверил скорость следующим образом:

import numpy as np

N = 1000
P = np.repeat(np.array([1, 2, 3, 4, 5]),N)
T = np.repeat(np.array([0, 0, 1, 1, 1]),N)

def using_bincount():
    U,inverse = np.unique(T,return_inverse=True)
    Q = np.bincount(inverse,weights=P)
    return Q,U
    # (array([  3.,  12.]), array([0, 1]))

def using_lc():
    U = list(set(T))
    Q = [sum([p for (p,t) in zip(P,T) if t == u]) for u in U]
    return Q,U

def using_slice():
    U = np.unique(T)
    Q = np.array([P[T == u].sum() for u in U])
    return Q,U

Для небольших массивов решение Вима быстрее (N = 1):

% python -mtimeit -s'import test' 'test.using_lc()'
100000 loops, best of 3: 18.4 usec per loop
% python -mtimeit -s'import test' 'test.using_slice()'
10000 loops, best of 3: 66.8 usec per loop
% python -mtimeit -s'import test' 'test.using_bincount()'
10000 loops, best of 3: 52.8 usec per loop

Для больших массивов решение Йориса быстрее (N = 1000):

% python -mtimeit -s'import test' 'test.using_lc()'
100 loops, best of 3: 9.93 msec per loop
% python -mtimeit -s'import test' 'test.using_slice()'
1000 loops, best of 3: 390 usec per loop
% python -mtimeit -s'import test' 'test.using_bincount()'
1000 loops, best of 3: 846 usec per loop

Я сомневаюсь, что это имеет значение в этом случае, но тесты могут меняться в зависимости от версии numpy, python, OS или аппаратного обеспечения. Не мешало бы повторить эти тесты на вашем компьютере.

Answer 2

import numpy as np
P = np.array([1, 2, 3, 4, 5]) 
T = np.array([0, 0, 1, 1, 1])

U = np.unique(T)
Q = np.array([P[T == u].sum() for u in U])

дает

In [17]: print Q, U
[3 12] [0 1]

Не совсем векторизовано, но быстрее, чем решение со списками.

Если вам нужны более мощные функции такого рода, возможно, вы можете взглянуть на панд .

Answer 3

>>> P = [1, 2, 3, 4, 5]; T = [0, 0, 1, 1, 1]
>>> U = list(set(T))
>>> Q = [sum([p for (p,t) in zip(P,T) if t == u]) for u in U]
>>> print Q, U
[3, 12] [0, 1]