Индекс одинаковых строк в массиве NumPy

Question

Я уже задавал вариант этого вопроса, но у меня все еще есть проблема, касающаяся времени выполнения моего кода.

Дан простой массив, состоящий из 15000 строк и 44 столбцов. Моя цель - выяснить, какие строки равны, и добавить их в список, например:

1 0 0 0 0
0 0 0 0 0 
0 0 0 0 0 
0 0 0 0 0 
1 0 0 0 0
1 2 3 4 5

Результат:

equal_rows1 = [1,2,3]
equal_rows2 = [0,4]

То, что я делал до сих пор, использует следующий код:

import numpy as np


input_data = np.load('IN.npy')

equal_inputs1 = []
equal_inputs2 = []

for i in range(len(input_data)):
  for j in range(i+1,len(input_data)):
     if np.array_equal(input_data[i],input_data[j]):
        equal_inputs1.append(i)
        equal_inputs2.append(j)

Проблема в том, что для возврата желаемых массивов требуется много времени, и это позволяет использовать только 2 разных «одинаковых списка строк», хотя их может быть больше. Есть ли лучшее решение для этого, особенно в отношении времени выполнения?

Answer 1

Вы можете использовать collections.defaultdict, который сохраняет значения строк в качестве ключей:

from collections import defaultdict

dd = defaultdict(list)

for idx, row in enumerate(df.values):
    dd[tuple(row)].append(idx)

print(list(dd.values()))
# [[0, 4], [1, 2, 3], [5]]

print(dd)
# defaultdict(<class 'list'>, {(1, 0, 0, 0, 0): [0, 4],
#                              (0, 0, 0, 0, 0): [1, 2, 3],
#                              (1, 2, 3, 4, 5): [5]})

Вы можете, если хотите, отфильтровать уникальные строки с помощью словарного понимания.

Answer 2

Это довольно просто с пандами groupby:

df
   A  B  C  D  E
0  1  0  0  0  0
1  0  0  0  0  0
2  0  0  0  0  0
3  0  0  0  0  0
4  1  0  0  0  0
5  1  2  3  4  5

[g.index.tolist() for _, g in df.groupby(df.columns.tolist()) if len(g.index) > 1]
# [[1, 2, 3], [0, 4]]

Если вы имеете дело со многими строками и множеством уникальных групп, это может быть немного медленным. Производительность зависит от ваших данных. Возможно, есть более быстрая альтернатива NumPy, но это, безусловно, проще всего понять.