您可以使用 matplotlib 手动完成。为了创建盒子,我们做了一个散点图,每种颜色都有正方形,但没有数据,所以它们不会出现。我们将这些散点图的返回值保存为句柄以传递给图例。然后我们从 heatmap 变量中获取 matplotlib figure 对象,该变量包含绘图所在的轴。在那里,我们使用自定义句柄和标签创建一个图例。
在该图上调用subplots_adjust,我们为右边的图例腾出空间。
import random
import numpy as np
import matplotlib
import seaborn as sb
import matplotlib.pyplot as plt
array = []
for x in range(10):
array.append(random.choices([-1,0, 1], k = 5))
array = np.array(array)
colors = ["red", "grey", "green"]
heatmap = sb.heatmap(array, cmap = ["red", "grey", "green"], cbar=False)
#Create dummy handles using scatter
handles = [plt.scatter([], [], marker='s', s=50, color=color) for color in colors]
labels = [-1, 0 , 1]
#Creating the legend using dummy handles
heatmap.figure.legend(handles=handles, labels=labels, loc='center right', frameon=False)
#Adjusting the plot space to the right to make room for the legend
heatmap.figure.subplots_adjust(right=0.8)
plt.show()
附注:
您可以用 numpy 函数替换生成随机数组的代码,这正是您想要的,但更方便。
所以替换这个:
array = []
for x in range(10):
array.append(random.choices([-1,0, 1], k = 5))
array = np.array(array)
有了这个:
array = np.random.choice((-1, 0, 1), (10, 5))
第一个参数是选项,第二个参数是数组的形状,所以在你的例子中是 10 x 5。