将 numpy 数组拆分为多个数组

Question

我有这样的x和y：

x = np.array([[2,3,5,6], [1,2,4,3], [1,5,6,4], [2,8,9,5]])
y = np.array([1,0,1,2])

我需要这样做。根据y值将数组分成多个数组。 X 和 y 数字可以不同，括号中的值的数量也可以不同。

x_0 = np.array([1,2,4,3])
x_1 = np.array([[2,3,5,6], [1,5,6,4]])
x_2 = np.array([2,8,9,5])

我尝试用字典和列表来做到这一点，但可能有人可以做得更短

Answer 1

这是满足您需要的一种方法：

result = [[] for _ in range(max(y) + 1)]
for idx, arr in zip(y, x):
    result[idx].append(arr)

如果您知道结果中恰好有三个列表，则：

x_0, x_1, x_2 = map(np.array, result)

结果：

>>> x_0
array([[1, 2, 4, 3]])
>>> x_1
array([[2, 3, 5, 6],
       [1, 5, 6, 4]])
>>> x_2
array([[2, 8, 9, 5]])

Answer 2

也许你可以使用字典理解：

x = np.array([[2, 3, 5, 6], [1, 2, 4, 3], [1, 5, 6, 4], [2, 8, 9, 5]])
y = np.array([1, 0, 1, 2])
separated_arrays = {value: x[y == value] for value in np.unique(y)}
print(separated_arrays)

输出：

{0: array([[1, 2, 4, 3]]), 
 1: array([[2, 3, 5, 6], [1, 5, 6, 4]]),
 2: array([[2, 8, 9, 5]])}