[英]Indices of unique values in n-dimensional array
我有一个包含 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]])]
我发现了这个相关的问题: Get a list of重复元素在 numpy array 中的所有索引,这可能会有所帮助,但我希望找到一个更直接的解决方案,不需要对数组进行展平和排序,并且尽可能高效.
这是一种矢量化方法,适用于任意维数的数组。 这个解决方案的想法是扩展return_index
方法的功能,并返回一个数组数组,每个数组包含一个 numpy 数组中唯一值的 N np.unique
。
对于更紧凑的解决方案,我定义了以下函数以及在不同步骤中的一些解释:
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)
您可以先在数组中获取非零元素,然后在列表理解中使用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]])]
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.