时间复杂度测试装饰器是否简洁合规？能否满足可视化分析需求

问题背景 #

我编写了一个装饰器，用于测试输入规模n增大时，特定操作或函数的时间复杂度特性（线性、二次等），代码如下：

实现代码 #

from typing import Union, Callable, Any, Tuple, List, Dict
import time
import numpy as np
import matplotlib.pyplot as plt

DataTypes = Union[List[int], Dict[int, int], None]

# pylint: disable=invalid-name
def data_factory(data_type: str, n: int) -> DataTypes:
    if data_type == "array":
        return list(range(n))
    if data_type == "dict":
        return {i: i for i in range(n)}
    if data_type is None:
        return None
    raise ValueError(f"Invalid data_type: {data_type}")


def time_complexity(
    data_type: str, repeat: int = 1, plot: bool = False
) -> Callable[[Callable[..., Any]], Callable[..., Tuple]]:
    def decorator(func: Callable[..., Any]) -> Callable[..., Tuple]:
        def wrapper(n_sizes: List[int], *args: Any, **kwargs: Dict[str, Any]) -> Tuple:
            avg_times = []
            median_times = []
            best_times = []
            worst_times = []

            for n in n_sizes:
                # create a list of n elements
                data_structure = data_factory(data_type, n)
                # note array is created outside the loop
                runtimes = []
                for _ in range(repeat):
                    start_time = time.perf_counter()

                    # pylint: disable=expression-not-assigned,line-too-long
                    func(n, data_structure, *args, **kwargs) if data_type else func(
                        n, *args, **kwargs
                    )  # <--- this is where it calls the function with n or data_structure as argument
                    end_time = time.perf_counter()
                    runtimes.append(end_time - start_time)

                avg_times.append(np.mean(runtimes))
                median_times.append(np.median(runtimes))
                best_times.append(np.min(runtimes))
                worst_times.append(np.max(runtimes))

            if plot:
                plt.figure(figsize=(10, 6))
                plt.plot(n_sizes, avg_times, "o-", label="Average")
                plt.plot(n_sizes, median_times, "o-", label="Median")
                plt.plot(n_sizes, best_times, "o-", label="Best")
                plt.plot(n_sizes, worst_times, "o-", label="Worst")
                plt.xlabel("Size of Input (n)")
                plt.ylabel("Execution Time (s)")
                plt.legend()
                plt.grid(True)
                plt.title(f"Time Complexity of {func.__name__}")
                plt.show()

            return n_sizes, avg_times, median_times, best_times, worst_times

        return wrapper

    return decorator

使用示例 #

@time_complexity(data_type="array", repeat=10, plot=True)
def list_append(n: int, array) -> None:
    array.append(n)


@time_complexity(data_type="array", repeat=10, plot=True)
def list_insert(n: int, array) -> None:
    array.insert(0, n)

# list_append(range(1000000, 10000001, 1000000))
# list_insert(range(1000000, 10000001, 1000000))

问题咨询 #

1. 排除CPU等不可控外部因素，该实现是否足以让我直观了解时间复杂度情况？
2. 这个装饰器是否简洁且功能符合预期？

回答 #

一、关于时间复杂度的直观评估有效性 #

整体框架具备核心价值，但存在关键逻辑bug会严重影响测试准确性，同时有可优化的细节：

1. 核心价值

通过多组n的运行时间统计（平均、中位数、最值）+ 可视化绘图，确实能直观反映时间随输入规模的增长趋势：比如常数复杂度是平线、线性是近似直线、二次是向上弯曲的曲线，足以帮你判断大致的复杂度类型。多次重复测试取统计值的设计，也能有效降低单次运行的随机性干扰。

2. 关键问题（必须修复）

当前代码中，针对每个n只生成一次数据结构，然后在repeat循环中反复修改它（比如list_append会不断增加数组长度，list_insert会不断在头部插入元素），导致每次重复测试的输入规模都不一样（第一次是n，第二次是n+1，以此类推），这会直接导致统计结果失真。

修复方案：把数据结构的生成放到repeat循环内部，保证每次测试都用全新的、规模为n的输入：

for n in n_sizes:
    runtimes = []
    for _ in range(repeat):
        # 每次重复都生成新的初始数据结构
        data_structure = data_factory(data_type, n)
        start_time = time.perf_counter()
        # 后续调用逻辑不变
        func(n, data_structure, *args, **kwargs) if data_type else func(n, *args, **kwargs)
        end_time = time.perf_counter()
        runtimes.append(end_time - start_time)

3. 其他优化点

• 对于 amortized 复杂度的操作（比如Python列表的append），可以增加repeat次数，更准确地体现平均性能。
• 可以添加对数坐标轴选项，方便观察非线性增长的趋势（比如O(log n)的曲线在对数x轴上是直线）。

二、关于装饰器的简洁性与功能符合性 #

1. 功能符合预期

装饰器的分层设计（三层结构传递data_type、repeat、plot参数）清晰，类型标注完整，覆盖了数据生成、多重复测试、多维度统计、可视化核心需求，基本达到了设计目标。

2. 简洁性优化空间

• 数据生成逻辑简化：把data_factory的条件判断改成字典映射，更简洁易维护：

  _DATA_FACTORIES = {
      "array": lambda n: list(range(n)),
      "dict": lambda n: {i: i for i in range(n)},
      None: lambda n: None
  }
  
  def data_factory(data_type: str, n: int) -> DataTypes:
      factory = _DATA_FACTORIES.get(data_type)
      if not factory:
          raise ValueError(f"Invalid data_type: {data_type}")
      return factory(n)
  

• 函数调用逻辑简化：替换冗长的三元表达式，让参数传递更清晰：

  # 替换原有的func调用代码
  call_args = [n, data_structure] + list(args) if data_type else [n] + list(args)
  func(*call_args, **kwargs)
  

• 统计代码简化：利用numpy数组一次性提取多维度统计值，减少重复代码：

  rt_array = np.array(runtimes)
  avg_times.append(rt_array.mean())
  median_times.append(rt_array.median())
  best_times.append(rt_array.min())
  worst_times.append(rt_array.max())
  

• 命名优化：data_type中的"array"对应Python的list，容易混淆，建议改成"list"更直观。
• 类型标注修正：**kwargs的正确类型标注应为**kwargs: Any，而非**kwargs: Dict[str, Any]。

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