Count unique elements row wise in an ndarray

Question

An extension to this question. In addition to having the unique elements row-wise, I want to have a similarly shaped array that gives me the count of unique values. For exam

Answer 1

This method does the same as np.unique for each row, by sorting each row and getting the length of consecutive equal values. This has complexity O(NMlog(M)) which is better than running unique on the whole array, since that has complexity O(NM(log(NM))

def row_unique_count(a):                                    
     args = np.argsort(a)
     unique = a[np.indices(a.shape)[0], args]
     changes = np.pad(unique[:, 1:] != unique[:, :-1], ((0, 0), (1, 0)), mode="constant", constant_values=1)
     idxs = np.nonzero(changes)
     tmp = np.hstack((idxs[-1], 0))
     counts = np.where(tmp[1:], np.diff(tmp), a.shape[-1]-tmp[:-1])
     count_array = np.zeros(a.shape, dtype="int")
     count_array[(idxs[0], args[idxs])] = counts
     return count_array

Running times:

In [162]: b = np.random.random(size=100000).reshape((100, 1000))

In [163]: %timeit row_unique_count(b)
100 loops, best of 3: 10.4 ms per loop

In [164]: %timeit count_unique_by_row(b)
100 loops, best of 3: 19.4 ms per loop

In [165]: assert np.all(row_unique_count(b) == count_unique_by_row(b))

Answer 2

The idea behind this answer is very similar to the one used here. I'm adding a unique imaginary number to each row. Therefore, no two numbers from different rows can be equal. Thus, you can find all the unique values in a 2D array per row with just one call to np.unique.

The index, ind, returned when return_index=True gives you the location of the first occurrence of each unique value.

The count, cnt, returned when return_counts=True gives you the count.

np.put(b, ind, cnt) places the count in the location of the first occurence of each unique value.

One obvious limitation of the trick used here is that the original array must have int or float dtype. It can not have a complex dtype to start with, since multiplying each row by a unique imaginary number may produce duplicate pairs from different rows.

---

import numpy as np

a = np.array([[1,  2, 2, 3,  4, 5],
              [1,  2, 3, 3,  4, 5],
              [1,  2, 3, 4,  4, 5],
              [1,  2, 3, 4,  5, 5],
              [1,  2, 3, 4,  5, 6]])

def count_unique_by_row(a):
    weight = 1j*np.linspace(0, a.shape[1], a.shape[0], endpoint=False)
    b = a + weight[:, np.newaxis]
    u, ind, cnt = np.unique(b, return_index=True, return_counts=True)
    b = np.zeros_like(a)
    np.put(b, ind, cnt)
    return b

yields

In [79]: count_unique_by_row(a)
Out[79]: 
array([[1, 2, 0, 1, 1, 1],
       [1, 1, 2, 0, 1, 1],
       [1, 1, 1, 2, 0, 1],
       [1, 1, 1, 1, 2, 0],
       [1, 1, 1, 1, 1, 1]])