基于范畴变量的numpy数组分割

Question

我试图根据分类变量“肥胖”来划分年龄和体重，然后用不同的颜色绘制两组。我想我可能把清单理解错了。当我绘图时，我只看到一种颜色和所有的数据点。在

import numpy as np 
import matplotlib.pyplot as plt
ages = np.array([20, 22, 23, 25, 27])
weights = np.array([140, 144, 150, 156, 160])
obese = np.array([0, 0, 0, 1, 1])

ages_normal = [ages for i in range(0, len(obese)) if obese[i] == 0]
weights_normal = [weights for i in range(0, len(obese)) if obese[i] == 0]

ages_obese = [ages for i in range(0, len(obese)) if obese[i] == 1]
weights_obese = [weights for i in range(0, len(obese)) if obese[i] == 1]

plt.scatter(ages_normal, weights_normal, color = "b")
plt.scatter(ages_obese, weights_obese, color = "r")
plt.show()

Answer 1

我可能会做一些类似的事情：

import numpy as np
import matplotlib.pyplot as plt
ages = np.array([20, 22, 23, 25, 27])
weights = np.array([140, 144, 150, 156, 160])
obese = np.array([0, 0, 0, 1, 1])

data = zip(ages, weights, obese)

data_normal = np.array([(a,w) for (a,w,o) in data if o == 0])
data_obese  = np.array([(a,w) for (a,w,o) in data if o == 1])

plt.scatter(data_normal[:,0], data_normal[:,1], color = "b")
plt.scatter(data_obese[:,0],  data_obese[:,1], color = "r")

plt.show()

但这可能更有效：

^{pr2}$

但你是对的，列表的理解有点偏差，也许你想要的是：

ages_normal = [ages[i] for i in range(0, len(obese)) if obese[i] == 0]
weights_normal = [weights[i] for i in range(0, len(obese)) if obese[i] == 0]

ages_obese = [ages[i] for i in range(0, len(obese)) if obese[i] == 1]
weights_obese = [weights[i] for i in range(0, len(obese)) if obese[i] == 1]

不同之处在于在ages/weights上添加了索引。

这三种方法都会生成您要查找的图形。