python - 在 numpy 数组的每一行中查找第一次出现的整数索引的最佳方法？

Question

如果我有一个数组，例如:

a = np.array([[1, 1, 2, 2, 1, 3, 4],
              [8, 7, 7, 7, 4, 8, 8]])

获得输出的最佳方式是什么:

array([[0, 2, 5, 6], [0, 1, 4]])

或

array([[0, 2, 5, 6], [4, 1, 0]])

这些是每行中每个整数第一次出现的索引。索引的顺序并不重要。

目前我正在使用:

res = []
for row in a:
  unique, unique_indexes = np.unique(a, return_index=True)
  res.append(unique_indexes)

但我想知道是否有 (num)pythonic 方法来避免 for 循环。

Answer 1

您可以以这样的方式转换数组，以便在一个批处理中处理整个事物。让我们从一个与您的问题非常相似的示例开始:

a = np.array([[1, 1, 2, 2, 1, 3, 4], [8, 7, 7, 7, 5, 8, 8]])

现在获取索引:

_, ix = np.unique(a, return_index=True)
# ix = array([ 0,  2,  5,  6, 11,  8,  7])

注意第一个元素的索引是正确的。以下元素偏移了 a 的大小。一般来说，偏移量是

offset = ix // a.shape[-1]
# offset = array([0, 0, 0, 0, 1, 1, 1])
ix %= a.shape[-1]
# ix = array([0, 2, 5, 6, 4, 1, 0])

您可以在 offset 更改值的每个位置对新的 ix 调用 np.split:

ix = np.split(ix, np.flatnonzero(np.diff(offset)) + 1)

那么为什么这个例子有效，而问题中的那个却无效呢？关键是 np.unique 使用基于排序的方法(这使得它在 O(n log n) 而不是 O(n) 的 collections.Counter)。这意味着为了使索引的顺序正确，每一行都必须与前一行不同并且大于前一行。请注意，在您的示例中，4 出现在两行中。您可以通过简单检查每一行中的 max 和 min 值来确保这一点:

mn = a.min(axis=1)
mx = a.max(axis=1)
diff = np.r_[0, (mx - mn + 1)[:-1].cumsum(0)] - mn
# diff = array([-4, -4])
b = a + diff[:, None]
# b = array([[0, 0, 1, 1, 0, 2, 3],
             [7, 6, 6, 6, 4, 7, 7]])

请注意，您必须将累积和偏移一个才能获得正确的索引。如果您处理大整数和/或非常大的数组，则在制作 diff 时必须更加小心以避免溢出。

现在您可以在调用 np.unique 时使用 b 代替 a。

长话短说

这是一个通用的无循环方法，适用于任何轴，而不仅仅是最后一个:

def global_unq(a, axis=-1):
    n = a.shape[axis]
    a = np.moveaxis(np.asanyarray(a), axis, -1).reshape(-1, n)
    mn = a.min(-1)
    mx = a.max(-1)
    diff = np.r_[0, (mx - mn + 1)[:-1].cumsum(0)] - mn
    _, ix = np.unique(a + diff[:, None], return_index=True)
    return np.split(ix % n, np.flatnonzero(np.diff(ix // n)) + 1)