使用索引拆分Numpy数组

Question

我有一个三维数组，内容是

pcar=[[xa ya za]
      [xb yb zb]
      .
      .
      [xn yn zn]]

还有一个索引数组，内容是

[0,1,0....,2]

这个索引数组告诉我，pcar中的每一行应该分到哪个聚类，这样我就可以用不同的颜色来绘制它们。聚类的最大值是3。

我想要的输出是

clusters[0] = [[xa ya za], [xc yc zc], ...]
clusters[1] = [[xb yb zb], ...]
clusters[2] = [..., [xn yn zn]]

然后我想把这些点在图上显示出来，每个聚类的点用不同的颜色表示。

Answer 1

类似这样的：

indices = [0,1,0....,2]
[pcar[i][index] for i, index in enumerate(indices)]
# [xa, yb, xc, ... zn]

Answer 2

根据我的理解，你想生成一个像下面这样的聚类列表：

clusters[0] = [[xa ya za], [xc yc zc], ...]
clusters[1] = [[xb yb zb], ...]
clusters[2] = [..., [xn yn zn]]

这用NumPy的索引功能来做是非常简单的。我们把索引数组叫做 indices。然后，

indices = np.asarray(indices)
num_clusters = 3
clusters = [pcar[indices==i] for i in xrange(num_clusters)]

就能得到你想要的结果。