3

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

例如，对于输入

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

我期待得到

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

我发现了这个相关问题： Get a list of all indices of repeat elements in a numpy array 这可能会有所帮助，但我希望找到一个更直接的解决方案，不需要对数组进行展平和排序，并且尽可能高效.

4

2 回答 2

4

这是一种矢量化方法，适用于任意维数的数组。此解决方案的想法是扩展 中return_index方法的功能np.unique，并返回一个数组数组，每个数组包含 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)

对于一维数组：

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)
于 2019-02-17T18:43:23.033 回答
3

您可以首先在数组中获取非零元素，然后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]])]
于 2019-02-17T15:17:18.577 回答