Быстрее return_inverse в np.unique

У меня есть большой массив 1D с более чем 100 миллионами элементов, и я применяю к нему np.unique

import numpy as np

x = np.random.randint(0,10000, size=100_000_000)

_, index = np.unique(x, return_inverse=True)

Что мне действительно нужно, так это index , возвращаемый из np.unique но мне вообще не нужен уникальный массив (т. np.unique Он одноразовый). Поскольку в моем реальном случае использования мне нужно np.unique вызывать np.unique для разных массивов (все с одинаковой длиной), это становится узким местом. Я предполагаю, что много времени уходит на сортировку уникального массива.

Какой самый быстрый способ получить index для большого одномерного массива (он может иметь длину более миллиарда элементов)?

Есть ли параллельный вариант?

Всего 1 ответ


Вот способ с использованием array-assignment + masking + indexing характерный для случая положительных целых чисел только во входном массиве x -

def return_inverse_only(x, maxnum=None):
    if maxnum is None:
        maxnum = x.max()+1 # Determines extent of indexing array
    p = np.zeros(maxnum, dtype=bool)
    p[x] = 1

    p2 = np.empty(maxnum, dtype=np.uint64)
    c = p.sum()
    p2[p] = np.arange(c)    
    out = p2[x]
    return out

Если максимальное число во входном массиве известно ранее, введите один добавленный номер как maxnum чтобы повысить производительность. в дальнейшем.

Сроки на больших массивах -

In [146]: np.random.seed(0)
     ...: x = np.random.randint(0,10000, size=100000)

In [147]: %timeit np.unique(x, return_inverse=True)
     ...: %timeit return_inverse_only(x)
     ...: %timeit return_inverse_only(x, maxnum=10000)
10.9 ms ± 229 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
539 µs ± 10.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
446 µs ± 30 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [148]: np.random.seed(0)
     ...: x = np.random.randint(0,10000, size=1000000)

In [149]: %timeit np.unique(x, return_inverse=True)
     ...: %timeit return_inverse_only(x)
     ...: %timeit return_inverse_only(x, maxnum=10000)
149 ms ± 5.92 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
6.1 ms ± 106 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
5.3 ms ± 504 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [150]: np.random.seed(0)
     ...: x = np.random.randint(0,10000, size=10000000)

In [151]: %timeit np.unique(x, return_inverse=True)
     ...: %timeit return_inverse_only(x)
     ...: %timeit return_inverse_only(x, maxnum=10000)
1.88 s ± 11.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
67.9 ms ± 1.66 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
55.8 ms ± 1.62 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

30x+ ускорение!


Есть идеи?

10000