matplotlib 3D散点图如何按条件为指定数据点设置颜色

核心修改逻辑：不要给散点图传入固定的单一颜色值，先基于判断规则为每个数据点生成对应的颜色值列表，再传入散点绘制接口即可实现分色效果。

对应需求梳理 #

• 绘制维度：共3个特征，分别为nbActionsD30、avgActionsMonth、actionSHR
• 着色规则：actionSHR值 >= 50的数据点使用专属高亮色，其余点使用默认蓝色
• 注意：原代码中%matplotlib widget属于IPython环境的魔法命令，需要放在Notebook文件的初始化导入位置，不要写在函数内部，避免执行异常。

修改后可直接运行的完整代码 #

# 魔法命令放在Notebook开头的导入区域，不要放入函数内
%matplotlib widget
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np

def plot3D(f1, f2, f3, data, threshold=50, highlight_color='#ff4757', default_color='blue'):
    xs = data[f1]
    ys = data[f2]
    zs = data[f3]

    # 基于阈值判断生成每个点对应的颜色
    point_colors = np.where(zs >= threshold, highlight_color, default_color)

    fig = plt.figure(figsize=(10,7))
    # 新版本matplotlib推荐的3D轴创建方式，兼容性更好，无版本警告
    ax = fig.add_subplot(111, projection='3d')

    ax.scatter(xs, ys, zs, s=50, c=point_colors, edgecolors="white")
    ax.set_xlabel(f1)
    ax.set_ylabel(f2)
    ax.set_zlabel(f3)

    # 可选：添加图例区分两类数据点，不需要可直接删除这段
    from matplotlib.lines import Line2D
    legend_handles = [
        Line2D([0], [0], marker='o', color='w', markerfacecolor=default_color, markersize=10, label=f'{f3} < {threshold}'),
        Line2D([0], [0], marker='o', color='w', markerfacecolor=highlight_color, markersize=10, label=f'{f3} >= {threshold}')
    ]
    ax.legend(handles=legend_handles)

    plt.show()

# 调用函数绘图
plot3D("avgActionsMonth", "nbActionsD30", "actionSHR", data)

关键参数说明 #

• 使用np.where做条件判断，生成和数据点长度完全一致的颜色数组，传给ax.scatter的c参数即可实现逐点着色，不要用传固定单一颜色的color参数
• 高亮色、默认色、判断阈值都已经设为函数可选参数，后续需要调整配色、修改判断阈值直接在调用时传参即可，不需要改动函数内部逻辑
• 原代码中直接实例化Axes3D(fig)的写法在matplotlib 3.4+版本会触发兼容性警告，替换为fig.add_subplot(111, projection='3d')的写法适配所有新版本。

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