Поиск в массиве всех совпадений и возвращение индексов совпадения

Question

Я пытаюсь создать массив, который содержит все строки, одна из которых (очень большой массив) совпадает с набором уникальных значений.Проблема состоит в том, что большой массив будет иметь несколько строк, где он будет совпадать, и мне нужно, чтобы все они хранились в одной строке нового массива.

Использование цикла for для обхода каждого из уникальных значений работает, но слишком медленно, чтобы его можно было использовать.Я искал векторизованное решение, но не увенчался успехом.Любая помощь будет принята с благодарностью!

    arrStart = []
    startRavel = startInforce['pol_id'].ravel()
    for policy in unique_policies:
        arrStart.append(np.argwhere(startRavel == policy))

Новый массив будет иметь ту же длину, что и массив уникальных значений, но каждый элемент будет списком всех строк, которые соответствуют этому уникальному значению в большом массиве..

Пример ввода будет выглядеть примерно так: startRavel = [1,2,2,2,3,3] unique_policies = [1,2,3]

Вывод: arrStart = [[0], [1,2,3], [4,5]]

Answer 1

Один из возможных вариантов с NumPy, аналогичный вашему, но сглаженный в понимании списка:

startRavel = np.array([1,2,2,2,3,3])
unique_policies = np.array([1,2,3])

[np.argwhere(startRavel == policy).flatten() for policy in unique_policies]
#=> [array([0]), array([1, 2, 3]), array([4, 5])]

---

Альтернатива, использующая flatnonzero():

[np.flatnonzero(startRavel == policy) for policy in unique_policies]

---

Версия генератора:

def matches_indexes(startRavel, unique_policies):
  for policy in unique_policies:
    yield np.flatnonzero(startRavel == policy)