Question

我有一个包含整数值的1D数组：

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

我想创建一个2D数组，第一个维度保存1D数组中整数值的索引：

idx = [np.where(a == (i+1)) for i in range(a.max())]

但这会返回一个列表（duh）：

type(idx)
list

第一个维度是一系列元组：

type(idx[0])
tuple

如何使用where子句从1D数组返回2D numpy值的索引数组？

编辑：

预期产出：

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

Answer 1

最接近2D阵列的是：

In [147]: np.array(tuple(np.where(a == e)[0] for e in np.unique(a)))
Out[147]: 
array([array([ 0, 14, 15, 16]),
       array([ 1,  4,  5,  7,  9, 10, 11, 13, 17, 19, 21]),
       array([ 2,  3,  6,  8, 12, 18, 20])], dtype=object)

但它是一维数组或数组。

你的部分问题是np.where返回一个数组元组，这样无论你的数组有多少维，它都会有相同的接口。由于你的只有一个，你可以得到0指数。

然后我会建议使用np.unique，因为它更好，但它会跳过a中不存在的索引。因此，如果这很重要，那么只需改回，但使用range(a.max() + 1)：

In [149]: np.array(tuple(np.where(a == e)[0] for e in range(a.max() + 1)))
Out[149]: 
array([array([], dtype=int64), array([ 0, 14, 15, 16]),
       array([ 1,  4,  5,  7,  9, 10, 11, 13, 17, 19, 21]),
       array([ 2,  3,  6,  8, 12, 18, 20])], dtype=object)

因为指数从0开始而不是1。

从1D数组numpy上的where子句创建2D数组

1 个答案: