Есть ли группа numpy по функциям?

Вдохновленный библиотекой Eelco Hoogendoorn, но без его библиотеки и с использованием того факта, что первый столбец вашего массива всегда увеличивается ( если нет, сначала выполните сортировку с помощью a = a[a[:, 0].argsort()])

>>> np.split(a[:,1], np.unique(a[:, 0], return_index=True)[1][1:])
[array([275, 441, 494, 593]),
 array([679, 533, 686]),
 array([559, 219, 455]),
 array([605, 468, 692, 613])]

Я не "timeit" ([EDIT] см. Ниже), но это, вероятно, более быстрый способ решить вопрос:

• Нет собственного цикла Python
• Списки результатов представляют собой массивы numpy, в случае, если вам нужно выполнить с ними другие операции numpy, новое преобразование не потребуется
• Сложность выглядит O (n) (с сортировкой она идет O (n log (n))

[ИЗМЕНИТЬ, сентябрь 2021 г.] Я запустил timeit на своем Macbook M1 для таблицы из 10 тыс. Случайных целых чисел. Продолжительность 1000 звонков.

>>> a = np.random.randint(5, size=(10000, 2))  # 5 different "groups"

# Only the sort
>>> a = a[a[:, 0].argsort()]
⏱ 116.9 ms

# Group by on the already sorted table
>>> np.split(a[:, 1], np.unique(a[:, 0], return_index=True)[1][1:])
⏱ 35.5 ms

# Total sort + groupby
>>> a = a[a[:, 0].argsort()]
>>> np.split(a[:, 1], np.unique(a[:, 0], return_index=True)[1][1:])
⏱ 153.0 ms 👑

# With numpy-indexed package (cf Eelco answer)
>>> npi.group_by(a[:, 0]).split(a[:, 1])
⏱ 353.3 ms

# With pandas (cf Piotr answer)
>>> df = pd.DataFrame(a, columns=["key", "val"]) # no timer for this line
>>> df.groupby("key").val.apply(pd.Series.tolist) 
⏱ 362.3 ms

# With defaultdict, the python native way (cf Piotr answer)
>>> d = defaultdict(list)
for key, val in a:
    d[key].append(val)
⏱ 3543.2 ms

# With numpy_groupies (cf Michael answer)
>>> aggregate(a[:,0], a[:,1], "array", fill_value=[])
⏱ 376.4 ms

Второй сценарий timeit, с 500 различными группами вместо 5. Я удивлен по поводу pandas, я запускал несколько раз, но в этом сценарии они просто плохо себя ведут.

>>> a = np.random.randint(500, size=(10000, 2))

just the sort  141.1 ms
already_sorted 392.0 ms
sort+groupby   542.4 ms
pandas        2695.8 ms
numpy-indexed  800.6 ms
defaultdict   3707.3 ms
numpy_groupies 836.7 ms

[EDIT] Я улучшил ответ благодаря ответ ns63sr и Behzad Shayegh (см. Комментарий) Также благодарим TMBailey за то, что вы заметили сложность argsort - n log (n).

Пакет numpy_indexed (отказ от ответственности: я являюсь его автором) призван заполнить этот пробел в numpy. Все операции в numpy-indexed полностью векторизованы, и ни один алгоритм O (n ^ 2) не пострадал во время создания этой библиотеки.

import numpy_indexed as npi
npi.group_by(a[:, 0]).split(a[:, 1])

Обратите внимание, что обычно более эффективно напрямую вычислять соответствующие свойства по таким группам (например, group_by (keys) .mean (values)), а не сначала разбивать их на список / массив с зубцами.

Numpy здесь не очень удобен, потому что желаемый результат не является массивом целых чисел (это массив объектов списка).

Я предлагаю либо чистый путь Python ...

from collections import defaultdict

%%timeit
d = defaultdict(list)
for key, val in a:
    d[key].append(val)
10.7 µs ± 156 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

# result:
defaultdict(list,
        {1: [275, 441, 494, 593],
         2: [679, 533, 686],
         3: [559, 219, 455],
         4: [605, 468, 692, 613]})

... или путь панд:

import pandas as pd

%%timeit
df = pd.DataFrame(a, columns=["key", "val"])
df.groupby("key").val.apply(pd.Series.tolist)
979 µs ± 3.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# result:
key
1    [275, 441, 494, 593]
2         [679, 533, 686]
3         [559, 219, 455]
4    [605, 468, 692, 613]
Name: val, dtype: object

n = np.unique(a[:,0])
np.array( [ list(a[a[:,0]==i,1]) for i in n] )

Выходы:

array([[275, 441, 494, 593], [679, 533, 686], [559, 219, 455],
       [605, 468, 692, 613]], dtype=object)

Упростив ответ Винсента Дж. и учитывая комментарий HS-туманности, можно использовать return_index = True вместо return_counts = True и избавьтесь от cumsum:

np.split(a[:,1], np.unique(a[:,0], return_index = True)[1])[1:]

Выход

[array([275, 441, 494, 593]),
 array([679, 533, 686]),
 array([559, 219, 455]),
 array([605, 468, 692, 613])]

Я использовал np.unique (), а затем np.extract ()

unique = np.unique(a[:, 0:1])
answer = []
for element in unique:
    present = a[:,0]==element
    answer.append(np.extract(present,a[:,-1]))
print (answer)

[array([275, 441, 494, 593]), array([679, 533, 686]), array([559, 219, 455]), array([605, 468, 692, 613])]

Учитывая X как массив элементов, которые вы хотите сгруппировать, и y (1D массив) как соответствующие группы, следующая функция выполняет группировку с помощью numpy :

def groupby(X, y):
    y = np.asarray(y)
    X = np.asarray(X)
    y_uniques = np.unique(y)
    return [X[y==yi] for yi in y_uniques]

Итак, groupby(a[:,1], a[:,0]) возвращает [array([275, 441, 494, 593]), array([679, 533, 686]), array([559, 219, 455]), array([605, 468, 692, 613])]

Мы также можем счесть полезным сгенерировать dict:

def groupby(X): 
    X = np.asarray(X) 
    x_uniques = np.unique(X) 
    return {xi:X[X==xi] for xi in x_uniques} 

Давайте попробуем:

X=[1,1,2,2,3,3,3,3,4,5,6,7,7,8,9,9,1,1,1]
groupby(X)                                                                                                      
Out[9]: 
{1: array([1, 1, 1, 1, 1]),
 2: array([2, 2]),
 3: array([3, 3, 3, 3]),
 4: array([4]),
 5: array([5]),
 6: array([6]),
 7: array([7, 7]),
 8: array([8]),
 9: array([9, 9])}

Обратите внимание, что это само по себе не очень привлекательно, но если мы сделаем X object или namedtuple, а затем предоставим функцию groupby, это станет более интересным. Я добавлю это позже.

Поздно на вечеринку, но все равно. Если вы планируете не только группировать массивы, но также хотите выполнять с ними операции, такие как сумма, среднее значение и т. Д., И вы делаете это с учетом скорости, вы также можете рассмотреть возможность использования numpy_groupies. Все эти групповые операции оптимизированы и связаны с numba. Они легко превосходят другие упомянутые решения.

from numpy_groupies.aggregate_numpy import aggregate
aggregate(a[:,0], a[:,1], "array", fill_value=[])
>>> array([array([], dtype=int64), array([275, 441, 494, 593]),
           array([679, 533, 686]), array([559, 219, 455]),
           array([605, 468, 692, 613])], dtype=object)
aggregate(a[:,0], a[:,1], "sum")
>>> array([   0, 1803, 1898, 1233, 2378])

Становится совершенно очевидным, что a = a[a[:, 0].argsort()] является узким местом всех конкурирующих алгоритмов группировки, во многом благодаря Винсенту Дж. для уточнения этого. Более 80% времени обработки просто тратится на этот метод argsort, и нет простого способа заменить или оптимизировать его. Пакет numba позволяет ускорить многие алгоритмы, и, надеюсь, argsort привлечет все усилия в будущем. оставшуюся часть группировки можно значительно улучшить, предполагая, что индексы первого столбца малы.

TL; DR

Остальная часть большинства методов группировки содержит метод np.unique, который является довольно медленным и избыточным в случаях, когда значения групп малы. Более эффективно заменить его на np.bincount, который позже можно будет улучшить в numba. Есть некоторые результаты того, как можно улучшить оставшуюся часть:

def _custom_return(unique_id, a, split_idx, return_groups):
    '''Choose if you want to also return unique ids'''
    if return_groups:
        return unique_id, np.split(a[:,1], split_idx)
    else: 
        return np.split(a[:,1], split_idx)

def numpy_groupby_index(a, return_groups=False):
    '''Code refactor of method of Vincent J'''
    u, idx = np.unique(a[:,0], return_index=True) 
    return _custom_return(u, a, idx[1:], return_groups)

def numpy_groupby_counts(a, return_groups=False):
    '''Use cumsum of counts instead of index'''
    u, counts = np.unique(a[:,0], return_counts=True)
    idx = np.cumsum(counts)
    return _custom_return(u, a, idx[:-1], return_groups)

def numpy_groupby_diff(a, return_groups=False):
    '''No use of any np.unique options'''
    u = np.unique(a[:,0])
    idx = np.flatnonzero(np.diff(a[:,0])) + 1
    return _custom_return(u, a, idx, return_groups)

def numpy_groupby_bins(a, return_groups=False):  
    '''Replace np.unique by np.bincount'''
    bins = np.bincount(a[:,0])
    nonzero_bins_idx = bins != 0
    nonzero_bins = bins[nonzero_bins_idx]
    idx = np.cumsum(nonzero_bins[:-1])
    return _custom_return(np.flatnonzero(nonzero_bins_idx), a, idx, return_groups)

def numba_groupby_bins(a, return_groups=False):  
    '''Replace np.bincount by numba_bincount'''
    bins = numba_bincount(a[:,0])
    nonzero_bins_idx = bins != 0
    nonzero_bins = bins[nonzero_bins_idx]
    idx = np.cumsum(nonzero_bins[:-1])
    return _custom_return(np.flatnonzero(nonzero_bins_idx), a, idx, return_groups)

Таким образом, numba_bincount работает так же, как np.bincount, и определяется следующим образом:

from numba import njit

@njit
def _numba_bincount(a, counts, m):
    for i in range(m):
        counts[a[i]] += 1

def numba_bincount(arr): #just a refactor of Python count
    M = np.max(arr)
    counts = np.zeros(M + 1, dtype=int)
    _numba_bincount(arr, counts, len(arr))
    return counts

Применение:

a = np.array([[1,275],[1,441],[1,494],[1,593],[2,679],[2,533],[2,686],[3,559],[3,219],[3,455],[4,605],[4,468],[4,692],[4,613]])
a = a[a[:, 0].argsort()]
>>> numpy_groupby_index(a, return_groups=False)
[array([275, 441, 494, 593]),
 array([679, 533, 686]),
 array([559, 219, 455]),
 array([605, 468, 692, 613])]
>>> numpy_groupby_index(a, return_groups=True)
(array([1, 2, 3, 4]),
 [array([275, 441, 494, 593]),
  array([679, 533, 686]),
  array([559, 219, 455]),
  array([605, 468, 692, 613])])

Тесты производительности

Для сортировки 100 миллионов элементов на моем компьютере (с 10 уникальными идентификаторами) требуется ~ 30 секунд. Проверим, сколько времени займут методы оставшейся части:

%matplotlib inline
benchit.setparams(rep=3)

sizes = [3*10**(i//2) if i%2 else 10**(i//2) for i in range(17)]
N = sizes[-1]
x1 = np.random.randint(0,10, size=N)
x2 = np.random.normal(loc=500, scale=200, size=N).astype(int)
a = np.transpose([x1, x2])

arr = a[a[:, 0].argsort()]
fns = [numpy_groupby_index, numpy_groupby_counts, numpy_groupby_diff, numpy_groupby_bins, numba_groupby_bins]
in_ = {s/1000000: (arr[:s], ) for s in sizes}
t = benchit.timings(fns, in_, multivar=True, input_name='Millions of events')
t.plot(logx=True, figsize=(12, 6), fontsize=14)

Без сомнения, bincount на основе numba — новый победитель наборов данных, содержащих небольшие идентификаторы. Это помогает сократить группировку отсортированных данных примерно в 5 раз, что составляет примерно 10% от общего времени выполнения.

Другой подход, предложенный Ashwini Chaudhary может быть тем, что вы ищете. Помещение его в простую функцию

def np_groupby(x, index):
    return np.split(x, np.where(np.diff(x[:,index]))[0]+1)

X = пустой массив

Индекс = индекс столбца

[0] + 1 согласно Ашвини, ... любое ненулевое значение означает, что элемент рядом с ним был другим, мы можем использовать numpy.where, чтобы найти индексы ненулевых элементов, а затем добавить 1, потому что фактический индекс такого элемента на единицу больше возвращаемого индекса; ...numpy.diff используется для определения фактического изменения элементов.