Return indices of non-NaN uniques in np.array of dtype=object

How do I return a list of indices corresponding to uniques in a np.array of dtype=object?

Analogous to:

arr = np.array(["one", "one", 2, 2])
result = np.unique(arr, return_inverse=True)[1]
print(result)
# [1, 1, 0, 0]

but NaN values included and those being ignored during indexing:

arr = np.array([nan, "one", 2, 2])
result = np.unique(arr, return_inverse=True)[1]
print(result)
# TypeError: '<' not supported between instances of 'float' and 'str'

I have already tried doing the following:

arr = np.array([nan, "one", 2, 2])
result = np.unique(arr[~np.isnan(arr)], return_inverse=True)[1]
print(result)
# TypeError: ufunc 'isnan' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the 

What I would like to get from the example above:

arr = np.array([nan, "one", 2, 2])
result = #...
print(result)
# [nan, 1, 0, 0]

Please note that arr is of dtype=object because it contains variable data types int and str.

Thank you in advance!


In the first example, the array is string dtype:

In [293]: arr = np.array(["one", "one", 2, 2])
In [294]: arr
Out[294]: array(['one', 'one', '2', '2'], dtype='<U21')
In [295]: np.unique(arr)
Out[295]: array(['2', 'one'], dtype='<U21')

If we specify object dtype

In [298]: arr = np.array(["one", "one", 2, 2], object)
In [299]: arr
Out[299]: array(['one', 'one', 2, 2], dtype=object)
In [300]: np.unique(arr)
Traceback (most recent call last):
  ...
  File "/usr/local/lib/python3.8/dist-packages/numpy/lib/arraysetops.py", line 333, in _unique1d
    ar.sort()
TypeError: '<' not supported between instances of 'int' and 'str'

Note the sort in the traceback.

What is your nan?

In [306]: arr = np.array([nan, "one", 2, 2])
Traceback (most recent call last):
  File "<ipython-input-306-abe4f4fe7b97>", line 1, in <module>
    arr = np.array([nan, "one", 2, 2])
NameError: name 'nan' is not defined

In [307]: arr = np.array([np.nan, "one", 2, 2])
In [308]: arr
Out[308]: array(['nan', 'one', '2', '2'], dtype='<U32')

nan is a float:

In [309]: arr = np.array([np.nan, 3, 2, 2])
In [310]: arr
Out[310]: array([nan,  3.,  2.,  2.])
In [311]: np.unique(arr)
Out[311]: array([ 2.,  3., nan])

unique on floats can be tricky, since floats aren't always "equal" if

np.unique uses np.lib.arraysetops._unique1d which has some special handling for nan (since nan isn't equal to anything, not even itself).

A sample string sorting

In [321]: np.sort(['one','a','B','','_',' '])
Out[321]: array(['', ' ', 'B', '_', 'a', 'one'], dtype='<U3')

It's been some time since I looked at string sorting (ASCII characters), so can't say exactly what the order is.