无需排序即可获得唯一的一维 NumPy 数组值

Question

I have many large 1D arrays and I'd like to grab the unique values. Typically, one could do:

x = np.random.randint(10000, size=100000000)
np.unique(x)

However, this performs an unnecessary sort of the array. The docs for np.unique do not mention any way to retrieve the indices without sorting. Other answers with np.unique include using return_index but, as I understand it, the array is still being sorted. So, I tried using set :

set(x)

But this is way slower than sorting the array with np.unique . Is there a faster way to retrieve the unique values for this array that avoids sorting and is faster than np.unique ?

Answer 1

If your values are positive integers in a relatively small range (eg 0 ... 10000), there is an alternative way to obtain a list of unique values using masks: (see unique2() below)

import numpy as np

def unique1(x):
    return np.unique(x)

def unique2(x):
    maxVal    = np.max(x)+1
    values    = np.arange(maxVal)
    used      = np.zeros(maxVal)
    used[x]   = 1
    return values[used==1]

# optimized (with option to provide known value range)
def unique3(x,maxVal=None):
    maxVal    = maxVal or np.max(x)+1
    used      = np.zeros(maxVal,dtype=np.uint8)
    used[x]   = 1
    return np.argwhere(used==1)[:,0]

In my tests this method is a lot faster than np.unique and it does not involve sorting:

from timeit import timeit
count = 3
x = np.random.randint(10000, size=100000000)

t = timeit(lambda:unique1(x),number=count)
print("unique1",t)

t = timeit(lambda:unique2(x),number=count)
print("unique2",t)

t = timeit(lambda:unique3(x),number=count)
print("unique3",t)

t = timeit(lambda:unique3(x,10000),number=count)
print("unique3",t, "with known value range")


# unique1 16.894681214000002
# unique2 0.8627655060000023
# unique3 0.8411087540000004
# unique3 0.5896318829999991 with known value range

Answer 2

Just in case you change your mind about dependencies, here's a dirt simple numba.njit implementation:

import numba

@numba.njit
def unique(arr):
    return np.array(list(set(arr)))


%timeit unique(x) #using Alain T.'s benchmark array
2.64 s ± 799 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit np.unique(x)
5.45 s ± 233 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Not as lightning fast as Above, but doesn't require positive integer inputs, either.