Numpy группа по нескольким векторам, получить групповые индексы

Question

У меня есть несколько массивов numpy;Я хочу создать групповой метод, который бы имел групповые идентификаторы для этих массивов.Затем он позволит мне индексировать эти массивы по идентификатору группы для выполнения операций над группами.

Например:

import numpy as np
import pandas as pd
a = np.array([1,1,1,2,2,3])
b = np.array([1,2,2,2,3,3])

def group_np(groupcols):
    groupby = np.array([''.join([str(b) for b in bs]) for bs in zip(*[c for c in groupcols])])
    _, groupby = np.unique(groupby, return_invesrse=True)
   return groupby

def group_pd(groupcols):
    df = pd.DataFrame(groupcols[0])
    for i in range(1, len(groupcols)):
        df[i] = groupcols[i]
    for i in range(len(groupcols)):
        df[i] = df[i].fillna(-1)
    return df.groupby(list(range(len(groupcols)))).grouper.group_info[0]

Выходы:

group_np([a,b]) -> [0, 1, 1, 2, 3, 4]
group_pd([a,b]) -> [0, 1, 1, 2, 3, 4]

Есть ли более эффективный способ его реализации, в идеале в чистом виде?Узким местом в настоящее время, кажется, является создание вектора, который будет иметь уникальные значения для каждой группы - в настоящий момент я делаю это путем объединения значений для каждого вектора в виде строк.

Я хочу, чтобы это работало для любого числавходные векторы, которые могут иметь миллионы элементов.

Редактировать: вот еще один тестовый пример:

a = np.array([1,2,1,1,1,2,3,1])
b = np.array([1,2,2,2,2,3,3,2])

Здесь все элементы группы 2,3,4,7 должны быть одинаковыми.

Edit2: добавление некоторых тестов.

a = np.random.randint(1, 1000, 30000000)
b = np.random.randint(1, 1000, 30000000)
c = np.random.randint(1, 1000, 30000000)

def group_np2(groupcols):
    _, groupby = np.unique(np.stack(groupcols), return_inverse=True, axis=1)
    return groupby

%timeit group_np2([a,b,c])
# 25.1 s +/- 1.06 s per loop (mean +/- std. dev. of 7 runs, 1 loop each)
%timeit group_pd([a,b,c])
# 21.7 s +/- 646 ms per loop (mean +/- std. dev. of 7 runs, 1 loop each)

Answer 1

Пакет numpy_indexed (dsiclaimer: я его автор) охватывает следующие варианты использования:

import numpy_indexed as npi
npi.group_by((a, b))

Передача подобного массива индексных массивов позволяет избежать создания копии;но если вы не против сделать копию, вы также можете использовать стек:

npi.group_by(np.stack(a, b))

Answer 2

После использования np.stack для массивов a и b, если для параметра return_inverse установлено значение True в np.unique, то это выход, который вы ищете:

a = np.array([1,2,1,1,1,2,3,1])
b = np.array([1,2,2,2,2,3,3,2])
_, inv = np.unique(np.stack([a,b]), axis=1, return_inverse=True)
print (inv)

array([0, 2, 1, 1, 1, 3, 4, 1], dtype=int64)

и вы можете заменить [a,b] в np.stack списком всех векторов.

Редактировать : используется более быстрое решение np.unique на sum массивов, умноженных на совокупное произведение (np.cumprod) max плюс 1 всех предыдущих массивов в groupcols.например:

def group_np_sum(groupcols):
    groupcols_max = np.cumprod([ar.max()+1 for ar in groupcols[:-1]])
    return np.unique( sum([groupcols[0]] +
                          [ ar*m for ar, m in zip(groupcols[1:],groupcols_max)]), 
                      return_inverse=True)[1]

Для проверки:

a = np.array([1,2,1,1,1,2,3,1])
b = np.array([1,2,2,2,2,3,3,2])
print (group_np_sum([a,b]))
array([0, 2, 1, 1, 1, 3, 4, 1], dtype=int64)

Примечание: число, связанное с каждой группой, может не совпадать (здесь я изменил первый элемент a на 3)

a = np.array([3,2,1,1,1,2,3,1])
b = np.array([1,2,2,2,2,3,3,2])
print(group_np2([a,b]))
print (group_np_sum([a,b]))
array([3, 1, 0, 0, 0, 2, 4, 0], dtype=int64)
array([0, 2, 1, 1, 1, 3, 4, 1], dtype=int64)

, но сами группы совпадают.

Теперь, чтобы проверить время:

a = np.random.randint(1, 100, 30000)
b = np.random.randint(1, 100, 30000)
c = np.random.randint(1, 100, 30000)
groupcols = [a,b,c]

%timeit group_pd(groupcols)
#13.7 ms ± 1.22 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit group_np2(groupcols)
#34.2 ms ± 6.88 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit group_np_sum(groupcols)
#3.63 ms ± 562 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)