如何在Plotly Express的px.scatter_3d()中展示KMeans聚类质心

解决方案 #

你可以直接在 Plotly Express 生成的聚类散点图基础上叠加质心图层，无需完全切换到底层 Plotly API 实现。首先修正原代码中px.scatter_3d缺失x/y/z坐标参数的问题，然后在指定位置添加质心绘制逻辑即可，完整可运行代码如下：

import numpy as np
import pandas as pd
from sklearn.cluster import KMeans
import plotly.express as px
import streamlit as st

df = pd.read_csv('iris.csv')
df_columns = ['column_a', 'column_b', 'column_c']

kmeans = KMeans(n_clusters=3, init = 'k-means++', max_iter=200)
km = kmeans.fit(df[df_columns])
centroids = km.cluster_centers_
cluster_labels = km.labels_
df['cluster'] = pd.Series(cluster_labels, index=df.index)

# 绘制原始样本聚类散点图，补全x/y/z参数
fig = px.scatter_3d(df, x='column_a', y='column_b', z='column_c', color=cluster_labels, labels={'color': 'cluster'})

# 叠加绘制聚类质心
fig.add_trace(
    px.scatter_3d(
        x=centroids[:, 0],
        y=centroids[:, 1],
        z=centroids[:, 2]
    ).update_traces(
        marker=dict(color='black', size=12, symbol='asterisk'),
        name='聚类质心'
    ).data[0]
)

st.plotly_chart(fig)

参数说明

• symbol='asterisk' 与你提供的Matplotlib示例中marker='*'效果完全一致
• 可自行调整size参数修改质心大小，保证质心比普通样本点更醒目
• 如需隐藏质心的图例项，可在update_traces中添加showlegend=False参数

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