np.unique はaxisを使うと遅い．

タイトルの通りで，numpyのuniqueはaxisを指定すると途端に遅くなります．
本記事では，2次元配列に絞って高速化する方法を模索します．

np.unique はいつ使うか？

NumPyのuniqueは以下のような場合に非常に強力です．

• 2次元配列のlexsortがしたい
• 2次元配列のlexsort orderをしりたい
• 重複除去をしたい

一方で，NumPyのuniqueはとても遅いことが有名であり，特にaxisを利用すると耐え難いほどに遅くなります．

まずはsortから．

Shapeが (N, M) なる配列 A において，lexsortをする方法は他にもあります．

例えば，lexsortを使えば，

order = lexsort(A[:, ::-1].T)

のように計算可能です．

また，argsortを利用する場合は

order = np.argsort(A[:, -1])
for m in range(1, M):
    order = order[np.argsort(A[order, -m-1], kind="stable")]

というように行えば，lexsortが可能です．
ここで，kind="stable"によって，sort時に同一値があれば，indexの若いものが前に来るようになります．

次に重複除去

sortさえできてしまえば，$i$ 番目要素と $i + 1$ 番目要素の比較のみを通して，重複判定可能です．
つまり，$i$ 番目要素と $i + 1$ 番目要素が異なるということは $i + 1$ 番目要素は少なくとも初出となります．
逆に，一致する場合は重複しているので除去すればよいです．

A_order = A[order]
is_first_occurrence = np.empty(order.size, dtype=bool)
is_first_occurrence[0] = True
is_first_occurrence[1:] = np.any(A_order[1:] != A_order[:-1], axis=-1)

A_uniq = A_order[is_first_occurrence]

ここで A_uniq が求めたかったuniqueかつlexsortされた配列です．

Inverseも求める．

return_inverse=True によって求めることのできる逆写像も求めてみます．

inv = np.empty(order.size, dtype=int)
inv[order] = np.cumsum(is_first_occurrence) - 1

ここで，A_uniq[inv] は A に一致します．

Indexも求める．

return_index=True によって求めることのできる配列も求めてみます．

inds = order[is_first_occurrence]

ここで，A[inds] は A_uniq に一致します．

Countも求める．

return_count=True によって求めることのできる配列も求めてみます．

first_inds = np.nonzero(is_first_occurrence)[0]
count = np.diff(np.append(first_inds, len(A)))

実行時間

import numpy as np

N, M = ..., 3
A = np.random.random((N, M))

def unique_custom(A: np.ndarray) -> np.ndarray:
    order = np.argsort(A[:, -1])
    for m in range(1, M):
        order = order[np.argsort(A[order, -m-1], kind="stable")]
    A_order = A[order]
    is_first_occurrence = np.empty(order.size, dtype=bool)
    is_first_occurrence[0] = True
    is_first_occurrence[1:] = np.any(A_order[1:] != A_order[:-1], axis=-1)
    return A_order[is_first_occurrence]

%timeit unique_custom(A)
%timeit np.unique(A, axis=0)

N	unique_custom	np.unique
$10^2$	13.7 μs ± 220 ns	33.3 μs ± 481 ns
$10^3$	88.5 μs ± 412 ns	284 μs ± 2.07 μs
$10^4$	1.49 ms ± 26.5 μs	3.63 ms ± 13.3 μs
$10^5$	19.9 ms ± 561 μs	45.6 ms ± 246 μs
$10^6$	263 ms ± 4.53 ms	519 ms ± 3.71 ms

上の表にあるとおりで，基本的に np.unique をそのまま使うよりも他の numpy メソッドを組み合わせるほうが速くなります．