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

Question

У меня есть большой numpy 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 для разных массивов (все с одинаковой длиной), это становится узким местом. Я предполагаю, что много времени тратится на сортировку уникального массива.

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

Есть ли параллельная опция?

Answer 1

Вот способ с array-assignment + masking + indexing указанием хитрости c для случая целых положительных чисел только во входном массиве 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

Если максимальное число во входном массиве известен ранее-hahnd, добавьте число с одним добавлением как maxnum, чтобы увеличить perf. далее.

Синхронизация больших массивов -

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+ ускорение!