如何获取二维数组中特定列的唯一行索引

Question

如果我有一个像这样的numpy索引……

import numpy as np

a = np.array([
    [0, 0],
    [0, 1],
    [1, 0],
    [1, 1],
])

我该如何找到在一个或多个指定列中值是唯一的行的索引呢？我的意思是……如果我指定某一列作为“掩码”，我该如何使用这列来找到唯一的行？比如说，如果我想要……

根据第0列（第0列是掩码）找唯一的行。我希望得到这样的结果……

[[0,1],[2,3]]

因为如果你用第0列作为唯一性的标准，行0和行1会被归为同一“唯一组”，而行2和行3会被归为另一个“唯一组”，因为它们在第0列的值是相同的。

如果我想根据第1列（现在第1列是掩码）来找行，我希望得到这样的输出……

[[0,2],[1,3]]

因为使用第1列作为唯一性的标准会导致行0和行2，以及行1和行3各自成为独立的唯一组，因为它们在第1列的值是相同的。

我还想能够根据多个列来获取唯一的行。所以如果我想根据第0列和第1列（现在第0列和第1列都是掩码）来找唯一的行，我希望得到这样的结果……

[[0],[1],[2],[3]]

因为当你同时使用这两列作为唯一性标准时，会有四个唯一的行。

在numpy中有没有简单的方法可以做到这一点？谢谢。

Answer 1

numpy_indexed 这个包（声明：我是它的作者）提供了一种完全向量化的解决方案来处理这类问题：

import numpy_indexed as npi
# entire rows of a determine uniqueness
npi.unique(a)
# only second column determines uniqueness
npi.unique(a[:, 1])

而且还可以处理更多类型的列。

Answer 2

一种可能的方法是使用循环：

import numpy

a = numpy.array([
    [0, 0],
    [0, 1],
    [1, 0],
    [1, 1],
])


un = numpy.unique(a)

results = []

# could be a list comprehension
for val in un:  

    # zero-th column, change as needed:   
    indices = a[:,0] == val  

    results.append(numpy.argwhere(indices).flatten())

result = numpy.array(results)

print result

根据你的需求和最终目标，你可以使用Pandas这个库。

它有一个叫做groupby的方法，你可以这样使用：

import pandas
import numpy as np

a = np.array([
    [0, 0],
    [0, 1],
    [1, 0],
    [1, 1],
])


df = pandas.DataFrame(a).groupby([0])  # zero-th column, change as needed

for key, group in df:
    print group.values

注意，这个方法返回的是实际的值，而不是索引。

Answer 3

试试用 itertools.groupby 这个工具

from itertools import groupby

data = [1,3,2,3,4,1,5,2,6,3,4]
data = [(x, k) for k, x in enumerate(data)]
data = sorted(data)

groups = []
for k, g in groupby(data, lambda x:x[0]):
    groups.append([x[1] for x in g])

print(groups)

输出结果是

[[0, 5], [2, 7], [1, 3, 9], [4, 10], [6], [8]]

Answer 4

这里有一个自定义的解决方案，虽然它的性能可能不太好，因为它做了很多复制操作，并且直接遍历了整个矩阵：

def groupby(a, key_columns):
    from collections import defaultdict
    groups = defaultdict(list)
    for i, row in enumerate(a):
        groups[tuple(row[c] for c in key_columns)].append(i)
    return groups.values()

这个方案假设 key_columns 是一个列表或元组，里面包含了你想要进行分组的相关列。你也可以检查一下参数，把单个索引变成一个只有一个元素的列表。

运行以下示例会得到这个输出：

>>> groupby(a, [0])
[[0, 1], [2, 3]]
>>> groupby(a, [1])
[[0, 2], [1, 3]]

它也适用于你提到的多个关键列：

>>> groupby(a, [0, 1])
[[1], [2], [0], [3]]

需要注意的是，由于使用了 defaultdict，所以值的顺序是没有保证的。你可以选择对结果值进行排序，或者根据你打算如何使用这些次级索引，使用 collections.OrderedDict。