如何将矩阵中的值与不等长项分组

Question

假设我有一个简单的数组：

a = np.arange(3)

以及具有相同长度的索引数组：

I = np.array([0, 0, 1])

我现在要根据索引对值进行分组。 如何对第一个数组的元素进行分组以生成下面的结果

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

---

以下是我尝试的：

a = np.arange(3)
I = np.array([0, 0, 1])
out = np.empty(2, dtype=object)
out.fill([])

aslists = np.vectorize(lambda x: [x], otypes=['object'])

out[I] += aslists(a)

但是，此方法不会串联列表，而是只保留每个索引的最后一个值：

array([[1], [2]], dtype=object)

---

或者，对于二维情况：

a = np.random.rand(100)
I = (np.random.random(100) * 5 //1).astype(int)
J = (np.random.random(100) * 5 //1).astype(int)

out = np.empty((5, 5), dtype=object)
out.fill([])

如何根据两个索引数组将项目从a追加到out

Answer 1

1D案例

假设I被排序，对于作为输出的数组列表-

idx = np.unique(I, return_index=True)[1]
out = np.split(a,idx)[1:]

另一个用slicing得到idx来分裂a-

out = np.split(a, np.flatnonzero(I[1:] != I[:-1])+1)

要获取列表数组作为输出-

np.array([i.tolist() for i in out])

样本运行-

In [84]: a = np.arange(3)

In [85]: I = np.array([0, 0, 1])

In [86]: out = np.split(a, np.flatnonzero(I[1:] != I[:-1])+1)

In [87]: out
Out[87]: [array([0, 1]), array([2])]

In [88]: np.array([i.tolist() for i in out])
Out[88]: array([[0, 1], [2]], dtype=object)

---

2D大小写

对于2D填充到2D数组的情况，我们可以这样做-

ncols = 5
lidx = I*ncols+J
sidx = lidx.argsort() # Use kind='mergesort' to keep order
lidx_sorted = lidx[sidx]
unq_idx, split_idx = np.unique(lidx_sorted, return_index=True)
out.flat[unq_idx] = np.split(a[sidx], split_idx)[1:]