Мы можем использовать тот факт, что элементы являются «простыми» (то есть неотрицательными и не слишком большими?) Целыми числами.
Хитрость заключается в том, чтобы создать разреженную матрицу с одним элементом в строке, а затем преобразоватьэто к столбцу мудрое представление.Это обычно быстрее, чем argsort
, потому что это преобразование O (M + N + nnz), если разреженная матрица MxN с ненулевыми значениями nnz.
from scipy import sparse
def use_sprsm():
x = sparse.csr_matrix((a, a, np.arange(a.size+1))).tocsc()
idx, = np.where(x.indptr[:-1] != x.indptr[1:])
return {i: a for i, a in zip(idx, np.split(x.indices, x.indptr[idx[1:]]))}
# for comparison
def use_asort():
idx = np.argsort(a)
el, c = np.unique(a, return_counts=True)
return dict(zip(el, np.split(idx, c.cumsum()[:-1])))
Пример выполнения:
>>> a = np.random.randint(0, 100, (10_000,))
>>>
# sanity check, note that `use_sprsm` returns sorted indices
>>> for k, v in use_asort().items():
... assert np.array_equal(np.sort(v), use_sprsm()[k])
...
>>> timeit(use_asort, number=1000)
0.8930604780325666
>>> timeit(use_sprsm, number=1000)
0.38419671391602606