如何为plotly.graph_objects.Heatmap设置离散色条？

我想用Plotly热力图可视化一组时空数据，以下是可运行的示例代码：

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import plotly.graph_objects as go

df = pd.DataFrame(d) # 使用对应的数据字典
n_locations = df.location_encoded.nunique()
pivot_table = df.pivot_table(index='day', columns='time', values='location_encoded', aggfunc='first')

heatmap_data = pivot_table.values

x_labels = pivot_table.columns
y_labels = pivot_table.index

cmap = plt.get_cmap('viridis', n_locations)

color_map = [cmap(i) for i in range(n_locations)]

fig = go.Figure(data=go.Heatmap(
    z=heatmap_data,
    x=x_labels,
    y=y_labels,
    colorscale=[[i / (n_locations - 1), f"rgba{color_map[i]}"] for i in range(n_locations)],
    colorbar=dict(
        tickvals=np.arange(n_locations),
        title='Location'
    ),
))

fig.update_layout(
    xaxis=dict(title='Time of Day'),
    yaxis=dict(title='Date'),
    title='Heatmap of Location Data',
)

fig.show()

这段代码生成的热力图右侧色条为连续刻度，但我需要将其改为离散刻度，即每个位置对应独立颜色。Plotly官网示例使用color_continuous_scale参数，但该参数仅适用于plotly.express可视化，请问如何为go.Heatmap设置离散色条？

更新：
以下代码可完美实现该需求（感谢评论区建议）：

# 首先创建一个np.linspace值列表，除首尾的0和1外，中间每个值重复两次
vals = np.r_[np.array(0), np.repeat(list(np.linspace(0, 1, self.n_locations+1))[1:-1], 2), np.array(1)]

# 然后生成包含值与对应颜色的列表
cc_scale = [[j, colors[i//2]] for i, j in enumerate(vals)]

# 创建Plotly热力图
self.fig = go.Figure(data=go.Heatmap(
    z=self.heatmap_data,
    x=x_labels,
    y=y_labels,
    colorscale=cc_scale,
    colorbar=dict(
        tickvals=np.linspace(1/self.n_locations/2, 1 - 1/self.n_locations/2, self.n_locations) * (self.n_locations - 1), # 让刻度居中
        ticktext=self.location_labels,
        title='Location'
    ),
))

内容的提问来源于stack exchange，提问作者sander