Question

我有一个像这样的numpy ndarray：

[
    [1,1,1],
    [2,2,2],
    [3,3,3],
    [4,4,4]
]

我想创建一个ndarray，其中包含所有4个内部数组，这些内部数组由以下类别数组[1,2,3,2]细分。我希望它看起来像这样：

[
    [
        [1,1,1]
    ], 
    [
        [2,2,2],[4,4,4]
    ],
    [
        [3,3,3]
    ]
]

有没有办法在python中优雅地做到这一点？

我之所以问这个问题，是因为我已经使用sklearns KMeans对数据集进行了聚类，现在我想要一种优雅的方式来处理每个类别的所有元素而又不会产生过多的循环。

Answer 1

您可以使用字典来映射数组及其各自的类别，并在最后收集输出。以这种方式使用字典的优点是它也能够处理出现顺序混乱的类别。

input_list = [
    [1,1,1],
    [2,2,2],
    [3,3,3],
    [4,4,4]
] 

category = [1, 2, 3, 2]

from collections import defaultdict

temp_dict = defaultdict(list)
#Iterate through each array of input list, and append to temp_dict
for i, arr in zip(category, input_list):
    temp_dict[i].append(arr)

result = list(temp_dict.values())
print(result)
#[[[1, 1, 1]], [[2, 2, 2], [4, 4, 4]], [[3, 3, 3]]]

（可选）如果顺序也很重要，请对字典键进行排序。

result = [x[1] for x in sorted(temp_dict.items())]
print(result)

在类别数组上细分ndarray

1 个答案: