Как посчитать выбор в (3, 2000) и быстрее?

Question

Есть ли способ ускорить следующие две строки кода?

choice = np.argmax(cust_profit, axis=0) 
taken = np.array([np.sum(choice == i) for i in range(n_pr)])

%timeit np.argmax(cust_profit, axis=0)
37.6 µs ± 222 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit np.array([np.sum(choice == i) for i in range(n_pr)])
40.2 µs ± 206 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

n_pr == 2
cust_profit.shape == (n_pr+1, 2000)

Решения:

%timeit np.unique(choice, return_counts=True)
53.7 µs ± 190 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit np.histogram(choice, bins=np.arange(n_pr + 2))
70.5 µs ± 205 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit np.bincount(choice)
7.4 µs ± 17.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Эти микросекунды меня беспокоят, потому что это код находится под двумя слоями scipy.optimize.minimize (method = 'Nelder-Mead'), который находится в двойном вложенном l oop, поэтому 40 мкс равняется 4 часам. И я думаю обернуть все это в genti c search.

Answer 1

Первая строка кажется довольно простой. Если вы не можете отсортировать данные или что-то в этом роде, вы застряли с линейным поиском в np.argmax. Вторую строку можно ускорить, просто используя numpy вместо vanilla python для ее реализации:

v, counts = np.unique(choice, return_counts=True)

В качестве альтернативы:

counts = np.histogram(choice, bins=np.arange(n_pr + 2))

Версия histogram optimized для целых чисел также существует:

count = np.bincount(choice)

Последние два варианта лучше, если вы хотите гарантировать, что корзины содержат все возможные значения choice, независимо от того, присутствуют ли они в массиве или нет.

При этом вам, вероятно, не стоит беспокоиться о том, что занимает микросекунды.