NumPy数组中相同行的索引

时间:2019-01-10 15:52:22

标签: python arrays pandas numpy

我已经问过这个问题的一个变种,但是我的代码运行时仍然有问题。

给出一个由15000行和44列组成的numpy数组。我的目标是找出哪些行相等,并将它们添加到列表中,如下所示:

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个不同的“相似行列表”。有没有更好的解决方案,尤其是在运行时方面?

2 个答案:

答案 0 :(得分:1)

这对于熊猫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替代方法,但这当然是最容易理解的。

答案 1 :(得分: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]})

您可以根据需要通过字典理解过滤出唯一的行。