1

Given some numpy array a

array([2,2,3,3,2,0,0,0,2,2,3,2,0,1,1,0])

what is the best way to get all groups of n indices with each of them having a different value in a?

Obviously there is no group larger than the number of unique elements in a, here 4.

So for example, one group of size 4 is

array([0,2,5,13])

Consider that a might be quite long, let's say up to 250k.

If the result gets too large, it might also be desirable not to compute all such groups, but only the first k requested.

Radio Controlled
  • 825
  • 8
  • 23

2 Answers2

1

For inputs as integers, we can have a solution based on this post -

In [41]: sidx = a.argsort() # use kind='mergesort' for first occurences

In [42]: c = np.bincount(a)

In [43]: np.sort(sidx[np.r_[0,(c[c!=0])[:-1].cumsum()]])
Out[43]: array([ 0,  2,  5, 13])

Another closely related to previous method for generic inputs -

In [44]: b = a[sidx]

In [45]: np.sort(sidx[np.r_[True,b[:-1]!=b[1:]]])
Out[45]: array([ 0,  2,  5, 13])

Another with numba for memory-efficiency and hence performance too, to select first indices along those unique groups and also with the additional k arg -

from numba import njit

@njit
def _numba1(a, notfound, out, k):
    iterID = 0
    for i,e in enumerate(a):
        if notfound[e]:
            notfound[e] = False
            out[iterID] = i
            iterID += 1
        if iterID>=k:
            break
    return out

def unique_elems(a, k, maxnum=None):
    # feed in max of the input array as maxnum value if known
    if maxnum is None:
        L = a.max()+1
    else:
        L = maxnum+1
    notfound = np.ones(L, dtype=bool)
    out = np.ones(k, dtype=a.dtype)
    return _numba1(a, notfound, out, k)

Sample run -

In [16]: np.random.seed(0)
    ...: a = np.random.randint(0,10,200)

In [17]: a
Out[17]: 
array([5, 0, 3, 3, 7, 9, 3, 5, 2, 4, 7, 6, 8, 8, 1, 6, 7, 7, 8, 1, 5, 9,
       8, 9, 4, 3, 0, 3, 5, 0, 2, 3, 8, 1, 3, 3, 3, 7, 0, 1, 9, 9, 0, 4,
       7, 3, 2, 7, 2, 0, 0, 4, 5, 5, 6, 8, 4, 1, 4, 9, 8, 1, 1, 7, 9, 9,
       3, 6, 7, 2, 0, 3, 5, 9, 4, 4, 6, 4, 4, 3, 4, 4, 8, 4, 3, 7, 5, 5,
       0, 1, 5, 9, 3, 0, 5, 0, 1, 2, 4, 2, 0, 3, 2, 0, 7, 5, 9, 0, 2, 7,
       2, 9, 2, 3, 3, 2, 3, 4, 1, 2, 9, 1, 4, 6, 8, 2, 3, 0, 0, 6, 0, 6,
       3, 3, 8, 8, 8, 2, 3, 2, 0, 8, 8, 3, 8, 2, 8, 4, 3, 0, 4, 3, 6, 9,
       8, 0, 8, 5, 9, 0, 9, 6, 5, 3, 1, 8, 0, 4, 9, 6, 5, 7, 8, 8, 9, 2,
       8, 6, 6, 9, 1, 6, 8, 8, 3, 2, 3, 6, 3, 6, 5, 7, 0, 8, 4, 6, 5, 8,
       2, 3])

In [19]: unique_elems(a, k=6)
Out[19]: array([0, 1, 2, 4, 5, 8])
Divakar
  • 218,885
  • 19
  • 262
  • 358
-1

Use Numpy.unique for this job. There are several other options, one can for instance return the number of times each unique item appears in a.

import numpy as np

# Sample data
a = np.array([2,2,3,3,2,0,0,0,2,2,3,2,0,1,1,0])

# The unique values are in 'u'
# The indices of the first occurence of the unique values are in 'indices'
u, indices = np.unique(a, return_index=True)
BramAppel
  • 1,346
  • 1
  • 9
  • 21