散点图 Matplotlib 的条件标记

Question

我试图将 x1 和 x2 相互绘制，如果 class == 1，那么标记应该是 +，如果 class == 0，那么标记应该只是一个点。

file.csv 的结构如下，作为条件的类列只能是 1 或 0:

x1	x2	班级
1	2	0
9	4	1
0	5	1
2	6	0

这是我的代码：

df = pd.read_csv('file.csv')
print(df)
x1 = df['x1'].to_numpy()
x2 = df['x2'].to_numpy()
class = df['class'].to_numpy()
        
figure, ax = plt.subplots()

for i in range(0,80,1):
    if label[i] > 0:
        ax.plot(x1[i], color="red", marker='o')
    else:
        ax.plot(x1[i], color="blue", marker='x')

结果应该是这样的：

Answer 1

我会以这种方式绘制所有数据，如果您要选择一个子集，请考虑使用新的数据框来选择子集：

import matplotlib.pyplot as plt

# the rows where class is 1 and plot them with '+' marker
plt.scatter(df[df['class']==1]['x1'], df[df['class']==1]['x2'], marker='+')

# the rows where class is 0 and plot them with 'o' marker 
plt.scatter(df[df['class']==0]['x1'], df[df['class']==0]['x2'], marker='o')

plt.xlabel('x1') 
plt.ylabel('x2') 
plt.title('Scatter plot x1 vs x2') 
plt.show()