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()

但这可能更有效：

data = np.array(np.vstack([ages, weights, obese])).T

ind_n = np.where(data[:,2] == 0)
ind_o = np.where(data[:,2] == 1)

plt.scatter(data[ind_n,0], data[ind_n,1], color = "b")
plt.scatter(data[ind_o,0], data[ind_o,1], color = "r")

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

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上添加了索引。

所有这三种方法都会生成您正在寻找的图表。

基于分类变量拆分numpy数组

1 个答案: