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

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

这是一个自定义解决方案，它肯定不会非常高效，因为它会进行大量复制并直接遍历矩阵:

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，具体取决于您计划如何使用二级索引。