Number of unique elements per row in a NumPy array

Question

For example, for

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

I want to get

[2, 2, 3]

Is there a way

Answer 1

This solution via np.apply_along_axis isn't vectorised and involves a Python-level loop. But it is relatively intuitive using len + np.unique functions.

import numpy as np
from toolz import compose

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

np.apply_along_axis(compose(len, np.unique), 1, a)    # [2, 2, 3]