a = np.random.randint(3, size=10)
a
index = np.arange(len(a))
sort_idx = np.argsort(a)
cnt = np.bincount(a)
np.split(index[sort_idx], np.cumsum(cnt[:-1]))
Найти индексы ненулевых элементов и сгруппировать по значениям