如何在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循环。

- aurelien_morel

---

由于您拥有的不是实际数组（不规则数组最好表示为列表），因此循环相当优化。 - Mad Physicist

如果你所说的“ragged”是指行的长度不同，那么实际上这里的初始数组a并不是“ragged”，我想。 - aurelien_morel

初始数组是的，但结果在一般情况下必然是不规则的。正如我所说，我找到了一个更好的解决方案。给我一分钟写出来。 - Mad Physicist

谢谢！是的，你说的结果是正确的。 - aurelien_morel

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])

您可以在新的ix中每次offset改变值的位置上调用np.split：

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

所以为什么这个示例是有效的，但问题中的那个示例不是呢？关键在于 np.unique 使用基于排序的方法（使其运行时间为 O(n log n) 而不是 collections.Counter 的 O(n)）。这意味着为了正确排序索引的顺序，每一行必须唯一且大于前一行。请注意，在你的示例中，4 出现在两行中。你可以通过检查每行中的最大值和最小值来确保这一点：

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]])

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

现在您可以在调用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)

- Mad Physicist

---

谢谢，它完美地运行了，有趣的方法！我不太了解numpy是否在复杂度方面比for循环有所改进，但无论如何都很好。 - aurelien_morel

@aurelien_morel。不存在这样的事情。对于排序较大的数组，会有一定的惩罚。 - Mad Physicist

Answer 2

你可以把它放到列表推导式中，但是你的循环已经相当干净了：

[list(np.unique(e, return_index=True)[-1]) for e in a]
# [[0, 2, 5, 6], [4, 1, 0]]

- Andreas