n维数组中唯一值的索引

Question

我有一个包含从0到n的值的2D Numpy数组。 我希望获得一个长度为n的列表，其中第i个元素是值为i + 1（不包括0）的所有索引的数组。

例如，对于输入：

array([[1, 0, 1],
   [2, 2, 0]])

我期望得到

[array([[0, 0], [0, 2]]), array([[1,0], [1,1]])]

我发现了这个相关问题： 如何在numpy数组中获取所有重复元素的索引列表 它可能会有所帮助，但我希望找到更直接的解决方案，而不需要对数组进行展平和排序，并且尽可能高效。

- kontradictos

---

你期望的结果是一个大小不同的数组列表。没有“直接”的方法。在拒绝它们之前，请仔细研究相关的答案。 - hpaulj

@hpaulj：正是我的观点。问题中要求的定制化太多了。而且期望的输出已经是一个列表了。 - Sheldore

你打算以“直接”和“高效”的方式使用生成的列表吗？ - hpaulj

参见：numpy.where的更快替代方案？ - Georgy

Answer 1

这里提供了一种向量化的方法，适用于具有任意维数的数组。这个解决方案的思路是扩展np.unique方法中return_index功能的功能，并返回一个包含NumPy数组中唯一值的N维索引数组的数组。为了更简洁地解决问题，我定义了以下函数，并在不同步骤中加上了一些解释。

def ndix_unique(x):
    """
    Returns an N-dimensional array of indices
    of the unique values in x
    ----------
    x: np.array
       Array with arbitrary dimensions
    Returns
    -------
    - 1D-array of sorted unique values
    - Array of arrays. Each array contains the indices where a
      given value in x is found
    """
    x_flat = x.ravel()
    ix_flat = np.argsort(x_flat)
    u, ix_u = np.unique(x_flat[ix_flat], return_index=True)
    ix_ndim = np.unravel_index(ix_flat, x.shape)
    ix_ndim = np.c_[ix_ndim] if x.ndim > 1 else ix_flat
    return u, np.split(ix_ndim, ix_u[1:])

---

检查来自问题的数组 -

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

vals, ixs = ndix_unique(a)

print(vals)
array([0, 1, 2])

print(ixs)
[array([[0, 1],
        [1, 2]]), 
 array([[0, 0],
        [0, 2]]), 
 array([[1, 0],
        [1, 1]])]

让我们尝试另一个案例：

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

vals, ixs = ndix_unique(a)

print(vals)
array([1, 2, 3, 4])

print(ixs)
array([array([[0, 0],
              [0, 1],
              [1, 2],
              [2, 2]]),
       array([[1, 0],
              [1, 1]]), 
       array([[2, 0],
              [2, 1]]),
       array([[0, 2]])], dtype=object)

对于一个1D数组：

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

vals, ixs = ndix_unique(a)

print(vals)
array([1, 3, 4, 5])

print(ixs)
array([array([0]), array([3, 4]), array([2]), array([1])], dtype=object)

最后，另一个带有3D ndarray的示例：

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

vals, ixs = ndix_unique(a)

print(vals)
array([1, 2, 3, 4])

print(ixs)
array([array([[0, 0, 0],
              [0, 0, 1]]),
       array([[0, 0, 2],
              [1, 0, 0]]), 
       array([[1, 0, 1]]),
       array([[1, 0, 2]])], dtype=object)

- yatu

---

啊，是的@Bazingaa，但请注意我需要同时使用y和其扁平化版本。因此，扁平化不是为了获取唯一值。 - yatu

有时候对于 unique 函数，使用 return_index 或 return_inverse 是很有用的。 - hpaulj

嗨@kontradictos。欢迎！我正在简化代码并更改一些内容。原因是它是从其他函数适应过来的，该函数考虑到唯一值向量（这里是x）无法排序。但在这里不是这种情况。几分钟后更新。 - yatu

@hpaulj 谢谢你的建议，顺便说一句，它有助于简化我的代码 :-) - yatu

Answer 2

你可以先获取数组中的非零元素，然后使用列表推导式中的argwhere函数获取每个非零元素的独立数组。在这里，np.unique(arr[arr!=0])将为您提供非零元素，您可以在其上迭代以获取索引。

arr = np.array([[1, 0, 1],
            [2, 2, 0]])

indices = [np.argwhere(arr==i) for i in np.unique(arr[arr!=0])]
# [array([[0, 0],
#         [0, 2]]), array([[1, 0],
#         [1, 1]])]

- Sheldore

---

这样我就不知道索引属于值1还是2。我想要得到两个索引列表，一个对应每个非零值。 - kontradictos

是的，这里与数组中不同的值没有对应关系。它只返回大于0的坐标，而不考虑它们的值。 - yatu

@yatu：请检查我的修改后的答案。感谢您的评论。 - Sheldore

现在是的，很难避免在这里使用for循环 :-) - yatu

你尝试将它应用到你的实际问题了吗？现在需要多长时间？ - Sheldore