为何绘制彩色分形曲线的Python代码运行速度如此缓慢？

问题根源 #

你当前代码速度慢的核心原因是循环调用plt.plot()绘制每一小段线段。Matplotlib每次调用plot()都会触发完整的渲染流程，当分形阶数提高后，线段数量呈指数增长（比如order=5时就有1024段），大量重复的渲染操作会导致性能急剧下降。

优化方案 #

使用Matplotlib的LineCollection组件批量绘制所有线段，同时用NumPy向量化操作计算颜色值，彻底避免循环带来的开销。

修改后的完整代码 #

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

#------------ plot stuff (i prefer dark mode) ------------------#

fig = plt.figure(figsize = (12,12),facecolor = "black")
ax = plt.axes()
ax.set_xlim(-1.1,1.1)
ax.set_ylim(-1.1,1.1)
ax.set_facecolor("black")
ax.spines["bottom"].set_color("white")
ax.spines["top"].set_color("white")
ax.spines["left"].set_color("white")
ax.spines["right"].set_color("white")
ax.tick_params(axis = "x", colors = "white") 
ax.tick_params(axis = "y", colors = "white")

#------------------------ fractal function ---------------------#

def fractal(order):
    def points(order):
        if (order==0):
            return np.array([-1+0j,1+0j])
        else:
            p1 = points(order-1)
            p2 = p1
            p1 = p1[:-1]
            p2 = p2[1:]
            dp = p2 - p1

            #---------------------  seed ---------------------#
            #--- this specific seed creates the koch curve ---#

            b = 4 # number of segments in the seed

            d1 = dp/3
            d2 = (1.5+np.sqrt(3)/2*1j)*dp/3
            d3 = 2*dp/3
            d4 = dp

            new_points = np.empty(len(p1)*(b+1), dtype=np.complex128)
            new_points[::b+1] = p1
            new_points[1::b+1] = p1 + d1
            new_points[2::b+1] = p1 + d2
            new_points[3::b+1] = p1 + d3
            new_points[4::b+1] = p1 + d4

            return new_points

    return_points = points(order)
    x, y = return_points.real, return_points.imag
    return x, y

#------------------ 优化后的绘图代码 ------------------#

x, y = fractal(order=5)
m = len(x)

# 1. 构造所有线段的数组格式：(N, 2, 2)，每个元素是线段的两个端点坐标
segments = np.stack([x[:-1], y[:-1], x[1:], y[1:]]).T.reshape(-1, 2, 2)

# 2. 向量化计算颜色值（Matplotlib接受0-1之间的RGB浮点数）
i = np.arange(m-1)
c = 255 / m
red = (255 + (2.5/m)*c*i*(i - m)) / 255  # 转成0-1范围
green = (0.5*c*i + 127) / 255
blue = 150 / 255

# 组合成颜色数组
colors = np.column_stack([red, green, blue])

# 3. 创建LineCollection并添加到坐标轴
line_collection = LineCollection(segments, colors=colors)
ax.add_collection(line_collection)

plt.show()

关键优化点说明 #

• LineCollection批量渲染：一次性将所有线段传递给Matplotlib，只触发一次渲染流程，性能比循环调用plot()提升几个数量级。
• 向量化颜色计算：用NumPy数组操作替代循环计算每个线段的颜色，避免Python循环的低效开销。
• 直接使用RGB浮点数：跳过十六进制颜色转换步骤，进一步减少计算量。

用这种方式，即使绘制order=10的科赫曲线（包含1048576段），也能在几秒内完成渲染，和无颜色版本的速度差距极小。

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