python - 根据分类变量分割numpy数组

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

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

最佳答案

我可能会做类似的事情：

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

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

关于python - 根据分类变量分割numpy数组，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/29176940/