CodeWars问题：优化照片计数Python代码至O(N)复杂度

问题分析 #

原代码的核心问题在于每次调用count(".")都会遍历子串，导致整体时间复杂度达到O(n²)，当输入字符串长度较大时会超时。我们需要将算法优化为严格O(n)的线性复杂度，才能通过所有测试用例。

优化思路 #

• 预处理前缀相机计数数组：prefix_cameras[i]存储从字符串开头到第i个位置（不含i）的相机总数
• 预处理后缀相机计数数组：suffix_cameras[i]存储从第i个位置（含i）到字符串末尾的相机总数
• 遍历字符串时，直接通过两个数组获取对应位置的相机数量，无需重复遍历子串，将单次统计操作从O(n)降为O(1)

优化后的代码 #

def count_photos(road):
    n = len(road)
    prefix_cameras = [0] * (n + 1)
    # 构建前缀相机计数
    for i in range(n):
        prefix_cameras[i+1] = prefix_cameras[i] + (1 if road[i] == "." else 0)
    
    suffix_cameras = [0] * (n + 1)
    # 构建后缀相机计数
    for i in range(n-1, -1, -1):
        suffix_cameras[i] = suffix_cameras[i+1] + (1 if road[i] == "." else 0)
    
    total = 0
    for idx, char in enumerate(road):
        if char == "<":
            # 向左行驶的车，统计左侧所有相机数
            total += prefix_cameras[idx]
        elif char == ">":
            # 向右行驶的车，统计右侧所有相机数
            total += suffix_cameras[idx+1]
    return total

代码解释 #

1. 前缀数组构建：从左到右遍历字符串，逐个累计相机数量，prefix_cameras[idx]等价于原代码中road[:idx].count(".")的结果
2. 后缀数组构建：从右到左遍历字符串，逐个累计相机数量，suffix_cameras[idx+1]等价于原代码中road[idx:].count(".")的结果
3. 最终统计：遍历字符串时，根据车辆行驶方向直接取对应数组的值累加，全程仅需三次线性遍历，时间复杂度严格为O(n)

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