Python中两个列表元素临近出现校验的实现及性能优化问题

你希望的for '1' in list1写法不符合Python语法规范，列表迭代只能逐个返回元素值，无法直接获取当前元素对应的索引，所以没法直接按你写的这个格式实现。不过可以用enumerate函数同时拿到索引和元素值，写出比原始版本更简洁的实现，性能也略优于你写的嵌套判断版本：

count = sum(
    1 for i, val in enumerate(list1)
    if val == '1' and '1' in list2[i:i+distance+1]
)

更高性能的实现方案 #

你原始版本的时间复杂度为O(n*d)，其中n为列表长度，d为设定的距离阈值，当d较大时重复遍历切片的开销会非常高。可以通过预处理+二分查找的方式把时间复杂度降到O(n + m log m)（m为list2中'1'的总数量），处理长列表时性能提升非常明显：

import bisect

def count_valid_ones(list1: list[str], list2: list[str], distance: int) -> int:
    # 提前收集list2所有'1'的位置，天然为升序排列
    ones_in_list2 = [idx for idx, val in enumerate(list2) if val == '1']
    count = 0
    list2_max_idx = len(list2) - 1
    for i, val in enumerate(list1):
        if val != '1':
            continue
        # 计算当前检查的右边界，避免越界
        right_bound = min(i + distance, list2_max_idx)
        # 二分查找第一个大于等于i的'1'的位置下标
        pos = bisect.bisect_left(ones_in_list2, i)
        # 若该位置存在且不超过右边界，说明符合要求
        if pos < len(ones_in_list2) and ones_in_list2[pos] <= right_bound:
            count += 1
    return count

如果你的列表规模特别大，还可以用Numpy的向量化操作进一步提速，适合批量处理的场景：

import numpy as np
from scipy.ndimage import maximum_filter1d

def count_valid_ones_np(list1: list[str], list2: list[str], distance: int) -> int:
    arr1 = np.array(list1, dtype=np.bool_)
    arr2 = np.array(list2, dtype=np.bool_)
    # 滑动窗口取最大值，窗口内有'1'则对应位置为True
    window_mask = maximum_filter1d(arr2, size=distance+1, mode='constant', cval=0, origin=-(distance//2))
    return int(np.sum(arr1 & window_mask))

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