根據另一個數組的值（未排序，但分組）將 NumPy 數組拆分為子數組

Question

假設我有兩個 NumPy arrays

x = [[1, 2, 8],
     [2, 9, 1],
     [3, 8, 9],
     [4, 3, 5],
     [5, 2, 3],
     [6, 4, 7],
     [7, 2, 3],
     [8, 2, 2],
     [9, 5, 3],
     [10, 2, 3],
     [11, 2, 4]]
y = [0, 0, 1, 0, 1, 1, 2, 2, 2, 0, 0] 

注意：（ x中的值沒有以任何方式排序。我選擇這個示例是為了更好地說明示例）（這些只是x和y的兩個示例x和y的值可以是任意多個不同的數字， y可以具有任意不同的數字，但x中的值總是與y中的值一樣多）

我想根據y中的值有效地將數組x拆分為子數組。

我想要的輸出是

z_0 = [[1, 2, 8],
       [2, 9, 1],
       [4, 3, 5],
       [10, 2, 3],
       [11, 2, 4]]
z_1 = [[3, 8, 9],
       [5, 2, 3],
       [6, 4, 7],]
z_2 = [[7, 2, 3],
       [8, 2, 2],
       [9, 5, 3]]

假設y從零開始並且沒有排序而是分組，那么最有效的方法是什么？

注意：這個問題是這個問題的未排序版本： Split a NumPy array into subarrays based on the values (sorted in up order) of another array

Answer 1

解決這個問題的一種方法是為每個y值建立一個過濾器索引列表，然后簡單地 select x的那些元素。 例如：

z_0 = x[[i for i, v in enumerate(y) if v == 0]]
z_1 = x[[i for i, v in enumerate(y) if v == 1]]
z_2 = x[[i for i, v in enumerate(y) if v == 2]]

Output

array([[ 1,  2,  8],
       [ 2,  9,  1],
       [ 4,  3,  5],
       [10,  2,  3],
       [11,  2,  4]])
array([[3, 8, 9],
       [5, 2, 3],
       [6, 4, 7]])
array([[7, 2, 3],
       [8, 2, 2],
       [9, 5, 3]])

如果您想更通用並支持y中的不同數字集，您可以使用理解來生成 arrays 的列表，例如

z = [x[[i for i, v in enumerate(y) if v == m]] for m in set(y)]

Output：

[array([[ 1,  2,  8],
       [ 2,  9,  1],
       [ 4,  3,  5],
       [10,  2,  3],
       [11,  2,  4]]),
 array([[3, 8, 9],
       [5, 2, 3],
       [6, 4, 7]]),
 array([[7, 2, 3],
       [8, 2, 2],
       [9, 5, 3]])]

如果y也是np.array並且與 x 長度相同，則可以簡化它以使用 boolean 索引：

z = [x[y==m] for m in set(y)]

Output 同上。

Answer 2

只需使用列表理解和 boolean 索引

x = np.array(x)
y = np.array(y)

z = [x[y == i] for i in range(y.max() + 1)]

z
Out[]: 
[array([[ 1,  2,  8],
        [ 2,  9,  1],
        [ 4,  3,  5],
        [10,  2,  3],
        [11,  2,  4]]),
 array([[3, 8, 9],
        [5, 2, 3],
        [6, 4, 7]]),
 array([[7, 2, 3],
        [8, 2, 2],
        [9, 5, 3]])]

Answer 3

略有變化。

from operator import itemgetter
label = itemgetter(1)

將隱含信息與 label... (index,label)相關聯

y1 = [thing for thing in enumerate(y)]

在 label 上排序

y1.sort(key=label)

按 label 分組並構造結果

import itertools
d = {}
for key,group in itertools.groupby(y1,label):
    d[f'z{key}'] = [x[i] for i,k in group]

---

Pandas解決方案：

>>> import pandas as pd
>>> >>> df = pd.DataFrame({'points':[thing for thing in x],'cat':y})
>>> z = df.groupby('cat').agg(list)
>>> z       
                                                points
cat
0    [[1, 2, 8], [2, 9, 1], [4, 3, 5], [10, 2, 3], ...
1                    [[3, 8, 9], [5, 2, 3], [6, 4, 7]]
2                    [[7, 2, 3], [8, 2, 2], [9, 5, 3]]