Применить функцию к массиву значений, используя уникальные элементы из другого

Question

Допустим, у меня есть два массива с одинаковыми размерами: массив индексов idx и массив значений val:

import numpy as np

idx = np.array([[10, 10, 13],
                [13, 10, 18],
                [10, 16, 18]])

np.random.seed(42)
val = np.round(np.random.uniform(0, 100, 9).reshape((3, 3)), 1)

print(repr(val))
# array([[30.5,  9.8, 68.4],
#        [44. , 12.2, 49.5],
#        [ 3.4, 90.9, 25.9]])

Как эффективно применить функцию, т. Е. mean к массиву значений, используя уникальные выборки из массива индекса?

Вот моя лучшая грубая попытка с for -циклом с ожидаемым результатом:

{i: val[idx == i].mean() for i in np.unique(idx)}
# {10: 13.975, 13: 56.2, 16: 90.9, 18: 37.7}

Хотя в этом примере это выглядит нормально, в реальном случае это занимает более 40 секунд с 6000 уникальными индексами в массиве с более чем 5M элементами. Это слишком долго, и я ищу более эффективный метод.

Answer 1

Для больших массивов Pandas должен оказаться более эффективным даже после учета затрат на настройку фрейма данных и преобразование серии в словарь:

import numpy as np, pandas as pd
# Python 3.7.0, Pandas 0.23.4, NumPy 1.15.1

np.random.seed(0)
n = 10**3
idx = np.random.randint(0, 20, (n, n))
val = np.random.random((n, n))

df = pd.DataFrame({'idx': idx.ravel(), 'val': val.ravel()})

%timeit pd.DataFrame({'idx': idx.ravel(), 'val': val.ravel()})  # 7.84 ms
%timeit df.groupby('idx')['val'].mean()                         # 34.6 ms
%timeit df.groupby('idx')['val'].mean().to_dict()               # 35.4 ms
%timeit {i: val[idx == i].mean() for i in np.unique(idx)}       # 108 ms