python - 如何根据第三列值用颜色绘制二维数据点

Question

假设我有一个包含三列的数组:

>>>print(arr)
array(
    [[1 2 -1]
    [1 3 1]
    [3 2 -1]
    [5 2 -1]]
)

假设我想把它变成散点图:

plt.plot(arr[:,0], arr[:,1], 'xr')

工作正常。但是，所有分散的点看起来都像一个红色“x”。假设如果第三列值为 -1，我希望有一个红色“x”，如果该值为 1，我希望有一个蓝色“o”。我该怎么做？

这可以通过plt.plot()实现吗？

Answer 1

如果要更改颜色和标记，则需要绘制多个散点图，每个标记至少绘制一个。

import matplotlib.pyplot as plt
import numpy as np

a = np.array([[1, 2, -1],
              [1, 3, 1],
              [3, 2, -1],
              [5, 2, -1]])
mapping= {-1: ("red", "x"), 1: ("blue", "o")}

for c in np.unique(a[:,2]):
    d = a[a[:,2] == c]
    plt.scatter(d[:,0], d[:,1], c=mapping[c][0], marker=mapping[c][1])
    #plt.plot(d[:,0], d[:,1], color=mapping[c][0], marker=mapping[c][1])

plt.show()

使用plot同样是可能的(请参阅代码中的注释行)。